Fixes in Imagenet training script (#1224)

Jaiaid · web-flow · commit a8483471779b · 2024-01-30T15:10:19.000-08:00
Fixes
* ngpus_per_node variable count fix
* dist_backend assertion added for nccl issue
diff --git a/imagenet/main.py b/imagenet/main.py
@@ -106,8 +106,12 @@ def main():
 
     if torch.cuda.is_available():
         ngpus_per_node = torch.cuda.device_count()
+        assert not (ngpus_per_node == 1 and args.dist_backend == "nccl"),\
+        "nccl backend requires GPU count>1, see https://github.com/NVIDIA/nccl/issues/103 perhaps use 'gloo'"
     else:
-        ngpus_per_node = 1
+        ngpus_per_node = 0
+        assert args.dist_backend != "nccl",\
+        "nccl backend does not work without GPU, see https://pytorch.org/docs/stable/distributed.html"
     if args.multiprocessing_distributed:
         # Since we have ngpus_per_node processes per node, the total world_size
         # needs to be adjusted accordingly
