glenn-jocher
commited on
Commit
•
7c2832c
1
Parent(s):
11121e3
assert equal model and dataset classes
Browse files
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)
|