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]>
- 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
|