Small bugfix to fix CI

mala-project · May 3, 2024 · b58c096 · b58c096
1 parent 873e486
commit b58c096
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mala/network/trainer.py b/mala/network/trainer.py
@@ -240,7 +240,9 @@ def _load_from_run(cls, params, network, data, file=None):
         # First, load the checkpoint.
         if params.use_ddp:
             map_location = {"cuda:%d" % 0: "cuda:%d" % get_local_rank()}
-        checkpoint = torch.load(file, map_location=map_location)
+            checkpoint = torch.load(file, map_location=map_location)
+        else:
+            checkpoint = torch.load(file)
 
         # Now, create the Trainer class with it.
         loaded_trainer = Trainer(