u

ACEsuit · Sep 10, 2024 · a5b0b12 · a5b0b12
1 parent f25be29
commit a5b0b12
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/mace/cli/run_train.py b/mace/cli/run_train.py
@@ -696,8 +696,11 @@ def run(args: argparse.Namespace) -> None:
             device=device,
         )
         model.to(device)
-        if args.distributed and args.device == "cuda":
-            distributed_model = DDP(model, device_ids=[local_rank])
+        if args.distributed:
+            if args.device == "cuda":
+                distributed_model = DDP(model, device_ids=[local_rank])
+            elif args.device == "cpu":
+                distributed_model = DDP(model, device_ids=[rank])
         model_to_evaluate = model if not args.distributed else distributed_model
         if swa_eval:
             logging.info(f"Loaded Stage two model from epoch {epoch} for evaluation")

diff --git a/mace/tools/slurm_distributed.py b/mace/tools/slurm_distributed.py
@@ -39,3 +39,4 @@ def __init__(self):
         self.world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
         self.local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
         self.rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
+        return