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
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)
|