Skip to content

Commit 3fbddc2

Browse files
authored
fix: Set correct device ID with TRT EP (#299)
1 parent f58118c commit 3fbddc2

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

src/onnxruntime.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -712,6 +712,11 @@ ModelState::LoadModel(
712712
values.push_back(value);
713713
}
714714
}
715+
716+
// assign correct GPU to EP
717+
keys.push_back(std::string("device_id"));
718+
values.push_back(std::to_string(instance_group_device_id));
719+
715720
std::vector<const char*> c_keys, c_values;
716721
if (!keys.empty() && !values.empty()) {
717722
for (size_t i = 0; i < keys.size(); ++i) {

0 commit comments

Comments
 (0)