Browse Source
Fix TRT `max_workspace_size` deprecation notice (#6856)
* Fix TRT `max_workspace_size` deprecation notice * Update export.py * Update export.pymodifyDataloader
Glenn Jocher
GitHub
2 years ago
1 changed files with 1 additions and 0 deletions
+ 1
- 0
export.py
View File
| 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) |
Loading…