From a1e1029e7ef30bfe0f09e32170bf8e0eb8fc946f Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 7 Mar 2022 13:52:53 +0100 Subject: [PATCH] Fix TRT `max_workspace_size` deprecation notice (#6856) * Fix TRT `max_workspace_size` deprecation notice * Update export.py * Update export.py --- export.py | 1 + 1 file changed, 1 insertion(+) diff --git a/export.py b/export.py index 7a5205d55ee6..1e3d3e2f2e71 100644 --- 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)