Merge pull request #602 from RandomDefaultUser/fix_ddp_training

Fixing DDP training
mala-project · Oct 30, 2024 · b3d117e · b3d117e
2 parents 27dec15 + 1881d10
commit b3d117e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mala/network/runner.py b/mala/network/runner.py
@@ -850,7 +850,12 @@ def _forward_entire_snapshot(
         # Ensure the Network is on the correct device.
         # This line is necessary because GPU acceleration may have been
         # activated AFTER loading a model.
-        self.network.to(self.network.params._configuration["device"])
+        if self.parameters_full.use_ddp:
+            self.network.module.to(
+                self.network.module.params._configuration["device"]
+            )
+        else:
+            self.network.to(self.network.params._configuration["device"])
 
         # Determine where the snapshot begins and ends.
         from_index = 0