Fix the bug that tensorRT batch_size does not take effect

Signed-off-by: wuhongsheng <[email protected]>
ultralytics · Jun 7, 2023 · 9f2b2a0 · 9f2b2a0
1 parent 89c3040
commit 9f2b2a0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarks.py b/benchmarks.py
@@ -76,7 +76,7 @@ def run(
             if f == '-':
                 w = weights  # PyTorch format
             else:
-                w = export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1]  # all others
+                w = export.run(weights=weights, imgsz=[imgsz], include=[f], batch_size=batch_size, device=device, half=half)[-1]  # all others
             assert suffix in str(w), 'export failed'
 
             # Validate