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:
1 parent
89c3040e73
commit
a199480ba6
1 file changed
+6
-1
+6
-1
@@ -76,7 +76,12 @@ def run(
|
|||||||
if f == '-':
|
if f == '-':
|
||||||
w = weights # PyTorch format
|
w = weights # PyTorch format
|
||||||
else:
|
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'
|
assert suffix in str(w), 'export failed'
|
||||||
|
|
||||||
# Validate
|
# Validate
|
||||||
|
|||||||
Reference in new issue
Block a user