Fixed issue with single gpu not having world_size

fhkiel-mlaip · Jul 14, 2020 · 787582f · 787582f
1 parent 6364892
commit 787582f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -449,9 +449,9 @@ def train(hyp, tb_writer, opt, device):
     opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size)))  # extend to 2 sizes (train, test)
     device = torch_utils.select_device(opt.device, apex=mixed_precision, batch_size=opt.batch_size)
     opt.total_batch_size = opt.batch_size
+    opt.world_size = 1
     if device.type == 'cpu':
         mixed_precision = False
-        opt.world_size = 1
     elif opt.local_rank != -1:
         # DDP mode
         assert torch.cuda.device_count() > opt.local_rank