From fa093f0f49b2d06e0295c593026f8f0ec5336d1e Mon Sep 17 00:00:00 2001 From: Kevin Chen Date: Fri, 29 Dec 2023 14:26:39 -0800 Subject: [PATCH] Update DDS check to also include previously used allocators --- .../providers/tensorrt/tensorrt_execution_provider.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 684303a8b6448..0272359dd3a03 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -937,13 +937,13 @@ Status BindContextOutput(Ort::KernelContext& ctx, // which we defer allocation until the size is known and don't call IExecution::setTensorAddress) // // Otherwise, if the shape of the output tensor is known prior to the runtime, ORT will pre-allocate memory buffer for the output tensor for enqueueV3. - if (is_dds_output) { - if (dds_output_allocator_map.find(output_name) == dds_output_allocator_map.end()) { + auto knownDDS = dds_output_allocator_map.find(output_name) != dds_output_allocator_map.end(); + if (is_dds_output || knownDDS) { + dds_output_set.emplace(output_name); + if (!knownDDS) { auto allocatorPtr = std::make_unique(); trt_context->setOutputAllocator(output_name, allocatorPtr.get()); dds_output_allocator_map[output_name] = std::move(allocatorPtr); - } else { - trt_context->setOutputAllocator(output_name, dds_output_allocator_map[output_name].get()); } } else { output_tensors[i] = ctx.GetOutput(output_index, output_shapes);