Revert "Fixes in Imagenet training script" (#1225)

pytorch · Jan 30, 2024 · ec8a172 · ec8a172
1 parent a848347
commit ec8a172
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/imagenet/main.py b/imagenet/main.py
@@ -106,12 +106,8 @@ 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 = 0
-        assert args.dist_backend != "nccl",\
-        "nccl backend does not work without GPU, see https://pytorch.org/docs/stable/distributed.html"
+        ngpus_per_node = 1
     if args.multiprocessing_distributed:
         # Since we have ngpus_per_node processes per node, the total world_size
         # needs to be adjusted accordingly