.fuse() after .eval()

This commit is contained in:
Glenn Jocher 2020-04-05 14:05:12 -07:00
parent d04738a27c
commit 968b2ec004
1 changed files with 3 additions and 3 deletions

View File

@ -34,12 +34,12 @@ def detect(save_img=False):
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights
modelc.to(device).eval() modelc.to(device).eval()
# Fuse Conv2d + BatchNorm2d layers
# model.fuse()
# Eval mode # Eval mode
model.to(device).eval() model.to(device).eval()
# Fuse Conv2d + BatchNorm2d layers
# model.fuse()
# Export mode # Export mode
if ONNX_EXPORT: if ONNX_EXPORT:
model.fuse() model.fuse()