DDP update

ultralytics · Aug 10, 2020 · 7f699d0 · 7f699d0
1 parent 657687c
commit 7f699d0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -51,7 +51,8 @@ def train(hyp, opt, device, tb_writer=None):
     init_seeds(2 + rank)
     with open(opt.data) as f:
         data_dict = yaml.load(f, Loader=yaml.FullLoader)  # model dict
-    check_dataset(data_dict)  # check
+    with torch_distributed_zero_first(rank):
+        check_dataset(data_dict)  # check
     train_path = data_dict['train']
     test_path = data_dict['val']
     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])  # number classes, names