move check_dataset() out of with context

ultralytics · Aug 10, 2020 · 648ed33 · 648ed33
1 parent f726f0c
commit 648ed33
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/test.py b/test.py
@@ -68,7 +68,7 @@ def test(data,
     model.eval()
     with open(data) as f:
         data = yaml.load(f, Loader=yaml.FullLoader)  # model dict
-        check_dataset(data)
+    check_dataset(data)  # check
     nc = 1 if single_cls else int(data['nc'])  # number of classes
     iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95
     niou = iouv.numel()

diff --git a/train.py b/train.py
@@ -51,7 +51,7 @@ 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_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