ONNX export explicit cpu map_location
This commit is contained in:
1 parent
ad71d2d513
commit
1e2cb6b480
1 file changed
+1
-1
@@ -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()
|
||||
|
||||
|
||||
Reference in new issue
Block a user