assert equal model and dataset classes

2020-06-04 17:21:22 -07:00 · 2020-06-04 17:21:22 -07:00 · 7c2832cd49
parent 11121e39ed
commit 7c2832cd49
1 changed files with 1 additions and 0 deletions
--- 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)