Detach all losses that get saved

nequip/train/trainer.py

@@ -574,6 +574,7 @@ def batch_step(self, data, n_batches, validation=False):
 
         self.model.train()
 
+        # Do any target rescaling
         data = data.to(self.device)
         data = AtomicData.to_AtomicDataDict(data)
         if hasattr(self.model, "unscale"):
@@ -583,6 +584,7 @@ def batch_step(self, data, n_batches, validation=False):
             # in train mode, if normalizes the targets
             data = self.model.unscale(data)
 
+        # Run model
         out = self.model(data)
 
         # If we're in evaluation mode (i.e. validation), then
@@ -593,54 +595,64 @@ def batch_step(self, data, n_batches, validation=False):
         loss, loss_contrib = self.loss(pred=out, ref=data)
 
         if not validation:
-
             self.optim.zero_grad()
             loss.backward()
             self.optim.step()
 
             if self.lr_scheduler_name == "CosineAnnealingWarmRestarts":
                 self.lr_sched.step(self.iepoch + self.ibatch / n_batches)
 
-        mae, mae_contrib = self.loss.mae(pred=out, ref=data)
-        scaled_loss_contrib = {}
-        if hasattr(self.model, "scale"):
-
-            for key in mae_contrib:
-                mae_contrib[key] = self.model.scale(
-                    mae_contrib[key], force_process=True, do_shift=False
-                )
+        # save loss stats
+        with torch.no_grad():
+            mae, mae_contrib = self.loss.mae(pred=out, ref=data)
+            scaled_loss_contrib = {}
+            if hasattr(self.model, "scale"):
 
-            # TO DO, this evetually needs to be removed. no guarantee that a loss is MSE
-            for key in loss_contrib:
+                for key in mae_contrib:
+                    mae_contrib[key] = self.model.scale(
+                        mae_contrib[key], force_process=True, do_shift=False
+                    )
 
-                scaled_loss_contrib[key] = {
-                    k: torch.clone(v) for k, v in loss_contrib[key].items()
-                }
+                # TO DO, this evetually needs to be removed. no guarantee that a loss is MSE
+                for key in loss_contrib:
 
-                scaled_loss_contrib[key] = self.model.scale(
-                    scaled_loss_contrib[key],
-                    force_process=True,
-                    do_shift=False,
-                    do_scale=True,
-                )
+                    scaled_loss_contrib[key] = {
+                        k: torch.clone(v) for k, v in loss_contrib[key].items()
+                    }
 
-                if "mse" in type(self.loss.funcs[key].func).__name__.lower():
                     scaled_loss_contrib[key] = self.model.scale(
                         scaled_loss_contrib[key],
                         force_process=True,
                         do_shift=False,
                         do_scale=True,
                     )
 
-        self.batch_loss = loss
-        self.batch_scaled_loss_contrib = scaled_loss_contrib
-        self.batch_loss_contrib = loss_contrib
-        self.batch_mae = mae
-        self.batch_mae_contrib = mae_contrib
+                    if "mse" in type(self.loss.funcs[key].func).__name__.lower():
+                        scaled_loss_contrib[key] = self.model.scale(
+                            scaled_loss_contrib[key],
+                            force_process=True,
+                            do_shift=False,
+                            do_scale=True,
+                        )
+
+            self.batch_loss = loss.detach()
+            self.batch_scaled_loss_contrib = {
+                k1: {k2: v2.detach() for k2, v2 in v1.items()}
+                for k1, v1 in scaled_loss_contrib.items()
+            }
+            self.batch_loss_contrib = {
+                k1: {k2: v2.detach() for k2, v2 in v1.items()}
+                for k1, v1 in loss_contrib.items()
+            }
+            self.batch_mae = mae.detach()
+            self.batch_mae_contrib = {
+                k1: {k2: v2.detach() for k2, v2 in v1.items()}
+                for k1, v1 in mae_contrib.items()
+            }
 
-        self.end_of_batch_log(validation)
-        for callback in self.end_of_batch_callbacks:
-            callback(self)
+            self.end_of_batch_log(validation)
+            for callback in self.end_of_batch_callbacks:
+                callback(self)
 
     @property
     def early_stop_cond(self):

scripts/train.py

@@ -39,7 +39,7 @@ def main():
 
     # Get statistics of training dataset
     (
-        (forces_std),
+        (forces_std,),
         (energies_mean, energies_std),
         (allowed_species, Z_count),
     ) = trainer.dataset_train.statistics(