* Fix TRT `max_workspace_size` deprecation notice * Update export.py * Update export.pymodifyDataloader
@@ -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) |