glenn-jocher commited on
Commit
acc58c1
1 Parent(s): a5a1760

Fix TRT `max_workspace_size` deprecation notice (#6856)

Browse files

* Fix TRT `max_workspace_size` deprecation notice

* Update export.py

* Update export.py

Files changed (1) hide show
  1. export.py +1 -0
export.py CHANGED
@@ -218,6 +218,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
218
  builder = trt.Builder(logger)
219
  config = builder.create_builder_config()
220
  config.max_workspace_size = workspace * 1 << 30
 
221
 
222
  flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
223
  network = builder.create_network(flag)
 
218
  builder = trt.Builder(logger)
219
  config = builder.create_builder_config()
220
  config.max_workspace_size = workspace * 1 << 30
221
+ # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice
222
 
223
  flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
224
  network = builder.create_network(flag)