Unverified 提交 acc58c1d authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Fix TRT `max_workspace_size` deprecation notice (#6856)

* Fix TRT `max_workspace_size` deprecation notice * Update export.py * Update export.py
上级 a5a1760e
...@@ -218,6 +218,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F ...@@ -218,6 +218,7 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
builder = trt.Builder(logger) builder = trt.Builder(logger)
config = builder.create_builder_config() config = builder.create_builder_config()
config.max_workspace_size = workspace * 1 << 30 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)) flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
network = builder.create_network(flag) network = builder.create_network(flag)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论