developer0hye glenn-jocher commited on
Commit
a1c3572
·
unverified ·
1 Parent(s): cf4f95b

Explicit `model.eval()` call `if opt.train=False` (#3475)

Browse files

* 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 <[email protected]>

Files changed (1) hide show
  1. models/export.py +1 -2
models/export.py CHANGED
@@ -58,8 +58,7 @@ if __name__ == '__main__':
58
  # Update model
59
  if opt.half:
60
  img, model = img.half(), model.half() # to FP16
61
- if opt.train:
62
- model.train() # training mode (no grid construction in Detect layer)
63
  for k, m in model.named_modules():
64
  m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
65
  if isinstance(m, models.common.Conv): # assign export-friendly activations
 
58
  # Update model
59
  if opt.half:
60
  img, model = img.half(), model.half() # to FP16
61
+ model.train() if opt.train else model.eval() # training mode = no Detect() layer grid construction
 
62
  for k, m in model.named_modules():
63
  m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
64
  if isinstance(m, models.common.Conv): # assign export-friendly activations