ONNX export explicit cpu map_location

5 years ago · 1e2cb6b480
parent ad71d2d513
commit 1e2cb6b480
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)['model']
+    model = torch.load(opt.weights, map_location=torch.device('cpu'))['model']
    model.eval()
    model.fuse()