Skip to content

Commit a1c3572

Browse files
Explicit model.eval() call if opt.train=False (#3475)
* call model.eval() when opt.train is False call model.eval() when opt.train is False * single-line if statement * cleanup Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
1 parent cf4f95b commit a1c3572

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

models/export.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,7 @@
5858
# Update model
5959
if opt.half:
6060
img, model = img.half(), model.half() # to FP16
61-
if opt.train:
62-
model.train() # training mode (no grid construction in Detect layer)
61+
model.train() if opt.train else model.eval() # training mode = no Detect() layer grid construction
6362
for k, m in model.named_modules():
6463
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
6564
if isinstance(m, models.common.Conv): # assign export-friendly activations

0 commit comments

Comments
 (0)