ONNX export explicit cpu map_location

ultralytics · Jun 12, 2020 · 1e2cb6b · 1e2cb6b
1 parent ad71d2d
commit 1e2cb6b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/onnx_export.py b/models/onnx_export.py
@@ -25,7 +25,7 @@
 
     # 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()