diff --git a/paddle/fluid/inference/tensorrt/engine.cc b/paddle/fluid/inference/tensorrt/engine.cc index c6c0d0479c9e5fc2c9139457fee43d0206c2af0d..0b3c099934c5714aaa90e8335484da4070983ce3 100644 --- a/paddle/fluid/inference/tensorrt/engine.cc +++ b/paddle/fluid/inference/tensorrt/engine.cc @@ -150,8 +150,12 @@ void TensorRTEngine::FreezeNetwork() { "Call InitNetwork first to initialize network.")); // build engine. infer_builder_->setMaxBatchSize(max_batch_); +#if IS_TRT_VERSION_GE(8300) + infer_builder_config_->setMemoryPoolLimit( + nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_); +#else infer_builder_config_->setMaxWorkspaceSize(max_workspace_); - +#endif bool enable_fp16 = (precision_ == AnalysisConfig::Precision::kHalf); if (enable_fp16) { bool support_fp16 = infer_builder_->platformHasFastFp16(); diff --git a/paddle/fluid/inference/tensorrt/test_tensorrt.cc b/paddle/fluid/inference/tensorrt/test_tensorrt.cc index 4a8bcb2565615eb7e63f2aa1e436acf98f475a15..571a2b56c56a385001e79f080d49c909903873ed 100644 --- a/paddle/fluid/inference/tensorrt/test_tensorrt.cc +++ b/paddle/fluid/inference/tensorrt/test_tensorrt.cc @@ -95,7 +95,11 @@ nvinfer1::IHostMemory* CreateNetwork() { network->markOutput(*output); // Build the engine. builder->setMaxBatchSize(1); +#if IS_TRT_VERSION_GE(8300) + config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 10); +#else config->setMaxWorkspaceSize(1 << 10); +#endif auto engine = builder->buildEngineWithConfig(*network, *config); EXPECT_NE(engine, nullptr); // Serialize the engine to create a model, then close.