ONNX export bug fix #93

5 years ago · e1e3399286
parent 8b26e89006
commit e1e3399286
1 changed files with 1 additions and 1 deletions
--- a/models/onnx_export.py
+++ b/models/onnx_export.py
@ -25,7 +25,7 @@ if __name__ == '__main__':

    # Load pytorch model
    google_utils.attempt_download(opt.weights)
-    model = torch.load(opt.weights, map_location=torch.device('cpu'))['model']
+    model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
    model.eval()
    model.fuse()