Divide args.workers by ngpus_per_node (pytorch#485)

RuiWang1998 · Dec 27, 2018 · c5985a8 · c5985a8
1 parent 44053c5
commit c5985a8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/imagenet/main.py b/imagenet/main.py
@@ -148,6 +148,7 @@ def main_worker(gpu, ngpus_per_node, args):
             # DistributedDataParallel, we need to divide the batch size
             # ourselves based on the total number of GPUs we have
             args.batch_size = int(args.batch_size / ngpus_per_node)
+            args.workers = int(args.workers / ngpus_per_node)
             model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
         else:
             model.cuda()