stride is an issue in the backward pass through the forces, this m mi…

…ght fix it
choderalab · Aug 24, 2024 · c5b87be · c5b87be
1 parent a315724
commit c5b87be
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/modelforge/train/training.py b/modelforge/train/training.py
@@ -463,7 +463,7 @@ def _get_forces(
 
         return {
             "per_atom_force_true": per_atom_force_true,
-            "per_atom_force_predict": per_atom_force_predict,
+            "per_atom_force_predict": per_atom_force_predict.contiguous(),
         }
 
     def _get_energies(
@@ -1128,6 +1128,10 @@ def setup_trainer(self) -> Trainer:
         """
         from lightning import Trainer
 
+        # if devices is a list
+        if isinstance(self.runtime_parameter.devices, list):
+            strategy = "ddp"
+
         trainer = Trainer(
             max_epochs=self.training_parameter.number_of_epochs,
             num_nodes=self.runtime_parameter.number_of_nodes,