deep_learning/ultralytics
1
0
Fork 0
mirror of https://github.com/ultralytics/ultralytics.git synced 2025-09-15 15:48:41 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Update

Browse Source
This commit is contained in:
Y-T-G 2025-09-14 23:48:51 +00:00
parent 0d2c5315ed
commit d72a07cf0e
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ultralytics/utils/export.py
View File

@ -100,9 +100,9 @@ def export_engine(
config = builder.create_builder_config()
workspace_bytes = int((workspace or 0) * (1 << 30))
is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10 # is TensorRT >= 10
if is_trt10 and workspace > 0:
if is_trt10 and workspace_bytes > 0:
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
elif workspace > 0: # TensorRT versions 7, 8
elif workspace_bytes > 0: # TensorRT versions 7, 8
config.max_workspace_size = workspace_bytes
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(flag)