saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-23 22:13:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[TensorRT EP] Fix bug for using correct device id for EP allocator (#17036)

Browse source 
The code always uses device id 0. Fix to use provider option
`device_id_`
This commit is contained in:
Chi Lo 2023-08-08 09:06:44 -07:00 committed by GitHub
parent 50719d2f8e
commit 5b9bf8b663
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -2272,7 +2272,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
std::unordered_set<std::string> input_names;
std::unordered_map<std::string, std::vector<int32_t>> tensor_shape_values;
OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, 0));
OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, device_id_), device_id_);
if (alloc_ == nullptr) {
Ort::ThrowOnError(api->KernelContext_GetAllocator(context, &mem_info, &alloc_));
}