include ability to do multigpu training

molmod · Jul 20, 2024 · b061b87 · b061b87
1 parent fdf2739
commit b061b87
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/psiflow/models/mace_utils.py b/psiflow/models/mace_utils.py
@@ -690,10 +690,12 @@ def main():
         type=str,
     )
     args = parser.parse_args()
-    if args.distributed:
-        world_size = torch.cuda.device_count()
+    world_size = torch.cuda.device_count()
+    if world_size > 1:
+        args.distributed = True
         import torch.multiprocessing as mp
 
         mp.spawn(run, args=(args, world_size), nprocs=world_size)
     else:
+        args.distributed = False
         run(0, args, 1)