From a199480ba6bb527598df11abbc1d679ccda82670 Mon Sep 17 00:00:00 2001 From: wuhongsheng <664116298@qq.com> Date: Thu, 8 Jun 2023 05:28:01 +0800 Subject: [PATCH] Fix the bug that tensorRT batch_size does not take effect (#11672) * Fix the bug that tensorRT batch_size does not take effect Signed-off-by: wuhongsheng <664116298@qq.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: wuhongsheng <664116298@qq.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- benchmarks.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/benchmarks.py b/benchmarks.py index fc3073965ab3..b590ff63cb01 100644 --- 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