From 7c2832cd4903e072bd4e20e175144101ff088313 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 4 Jun 2020 17:21:22 -0700 Subject: [PATCH] assert equal model and dataset classes --- train.py | 1 + 1 file changed, 1 insertion(+) diff --git a/train.py b/train.py index 6552a6e..f963391 100644 --- 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)