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>
This commit is contained in:
wuhongshengandpre-commit-ci[bot] authored and GitHub committed 2023-06-07 23:28:01 +02:00
1 parent 89c3040e73
commit a199480ba6
1 file changed
+6 -1
+6 -1
View File
@@ -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