From d72a07cf0e92383eec550671aeb9019ca21a0e50 Mon Sep 17 00:00:00 2001 From: Y-T-G <32206511+Y-T-G@users.noreply.github.com> Date: Sun, 14 Sep 2025 23:48:51 +0000 Subject: [PATCH] Update --- ultralytics/utils/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ultralytics/utils/export.py b/ultralytics/utils/export.py index d58022f7c8..3c1dde9fd2 100644 --- a/ultralytics/utils/export.py +++ b/ultralytics/utils/export.py @@ -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)