uuguuguguguuuu

pytorch · Jul 11, 2023 · 06ab751 · 06ab751
1 parent f956d01
commit 06ab751
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/references/detection/utils.py b/references/detection/utils.py
@@ -262,9 +262,9 @@ def init_distributed_mode(args):
         args.rank = int(os.environ["RANK"])
         args.world_size = int(os.environ["WORLD_SIZE"])
         args.gpu = int(os.environ["LOCAL_RANK"])
-    # elif "SLURM_PROCID" in os.environ:
-    #     args.rank = int(os.environ["SLURM_PROCID"])
-    #     args.gpu = args.rank % torch.cuda.device_count()
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
     else:
         print("Not using distributed mode")
         args.distributed = False