From 1e2cb6b48014f38be8337a976cb2fcc177edf898 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 12 Jun 2020 13:54:00 -0700
Subject: [PATCH] ONNX export explicit cpu map_location

---
 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 a7132de8465..fd8abeb6365 100644
--- 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()