Fix TRT max_workspace_size deprecation notice (#6856)

* Fix TRT `max_workspace_size` deprecation notice * Update export.py * Update export.py
ultralytics · Mar 7, 2022 · acc58c1 · acc58c1
1 parent a5a1760
commit acc58c1
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/export.py b/export.py
@@ -218,6 +218,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
         builder = trt.Builder(logger)
         config = builder.create_builder_config()
         config.max_workspace_size = workspace * 1 << 30
+        # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)  # fix TRT 8.4 deprecation notice
 
         flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
         network = builder.create_network(flag)