From 08ddc475df060054608beded195c3a829cdec756 Mon Sep 17 00:00:00 2001
From: milesial <milesial@users.noreply.github.com>
Date: Fri, 7 Mar 2025 17:27:30 -0800
Subject: [PATCH] Set correct device ID with TRT EP

---
 src/onnxruntime.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/onnxruntime.cc b/src/onnxruntime.cc
index 6eae3c4..40301da 100644
--- a/src/onnxruntime.cc
+++ b/src/onnxruntime.cc
@@ -712,6 +712,11 @@ ModelState::LoadModel(
                     values.push_back(value);
                   }
                 }
+
+                // assign correct GPU to EP
+                keys.push_back(std::string("device_id"));
+                values.push_back(std::to_string(instance_group_device_id));
+
                 std::vector<const char*> c_keys, c_values;
                 if (!keys.empty() && !values.empty()) {
                   for (size_t i = 0; i < keys.size(); ++i) {