glenn-jocher commited on
Commit
7c2832c
1 Parent(s): 11121e3

assert equal model and dataset classes

Browse files
Files changed (1) hide show
  1. train.py +1 -0
train.py CHANGED
@@ -76,6 +76,7 @@ def train(hyp):
76
 
77
  # Create model
78
  model = Model(opt.cfg).to(device)
 
79
 
80
  # Image sizes
81
  gs = int(max(model.stride)) # grid size (max stride)
 
76
 
77
  # Create model
78
  model = Model(opt.cfg).to(device)
79
+ assert model.md['nc'] == nc, '%s nc=%g classes but %s nc=%g classes' % (opt.data, nc, opt.cfg, model.md['nc'])
80
 
81
  # Image sizes
82
  gs = int(max(model.stride)) # grid size (max stride)