Merge branch 'master' into snyk-fix-19a9bd869ca677b68dcdaf5f4affcd24

2024-07-25 14:03:26 +02:00 · 2024-07-25 14:03:26 +02:00 · 96d3c360a2
parent d7173ca923 6deb2d75cb
commit 96d3c360a2
1 changed files with 1 additions and 1 deletions
--- 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)