From e1e33992865ddc9eb5a0d7f0ed1ccfb7d7a33e8c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 16 Jun 2020 09:59:42 -0700 Subject: [PATCH] ONNX export bug fix #93 --- models/onnx_export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/onnx_export.py b/models/onnx_export.py index fd8abeb..59d13b1 100644 --- 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()