assert equal model and dataset classes

ultralytics · Jun 5, 2020 · 7c2832c · 7c2832c
1 parent 11121e3
commit 7c2832c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -76,6 +76,7 @@ def train(hyp):
 
     # Create model
     model = Model(opt.cfg).to(device)
+    assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc'])
 
     # Image sizes
     gs = int(max(model.stride))  # grid size (max stride)