Convert BatchNorm to SyncBatchNorm

ultralytics · Jul 6, 2020 · 77c8e27 · 77c8e27
1 parent 96fa40a
commit 77c8e27
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -194,6 +194,7 @@ def train(hyp, tb_writer, opt, device):
     # DDP mode
     if device.type != 'cpu' and opt.local_rank != -1:
         # pip install torch==1.4.0+cku100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
         if mixed_precision:
             model = DDP(model, delay_allreduce=True)
         else: