Merge pull request #48 from ultralytics/master

Fix the bug that tensorRT batch_size does not take effect (ultralytics#11672)
NagatoYuki0943 · Jun 9, 2023 · c8da105 · c8da105
2 parents 6a87357 + a199480
commit c8da105
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/benchmarks.py b/benchmarks.py
@@ -76,7 +76,12 @@ 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