diff --git a/export.py b/export.py
index 30944a93aff8..f5db08096025 100644
--- a/export.py
+++ b/export.py
@@ -256,7 +256,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:"
         # Load model
         weights = 'yolov5s.pt'
         device = select_device('')
-        model = attempt_load(weights, map_location=device)
+        model = attempt_load(weights, device=device)
 
         # Example input tensor
         im = torch.zeros(1, 3, 640, 640).to(device)