ONNX export explicit cpu map_location

This commit is contained in:
Glenn Jocher committed 2020-06-12 13:54:00 -07:00
1 parent ad71d2d513
commit 1e2cb6b480
1 file changed
+1 -1
+1 -1
View File
@@ -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()