saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-03 03:58:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Python with DmlExecutionProvider : choose device_id in SessionOptions (#7964)

Browse source
This commit is contained in:
iperov 2021-06-09 00:02:12 +04:00 committed by GitHub
parent 4ecbae43b2
commit 66170bfcef
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
onnxruntime/python/onnxruntime_pybind_state.cc
View file

@ -622,7 +622,18 @@ static void RegisterExecutionProviders(InferenceSession* sess, const std::vector
#endif
} else if (type == kDmlExecutionProvider) {
#ifdef USE_DML
RegisterExecutionProvider(sess, *onnxruntime::CreateExecutionProviderFactory_DML(0));
int device_id = 0;
auto it = provider_options_map.find(type);
if (it != provider_options_map.end()) {
for (auto option : it->second) {
if (option.first == "device_id") {
if (!option.second.empty()) {
device_id = std::stoi(option.second);
}
}
}
}
RegisterExecutionProvider(sess, *onnxruntime::CreateExecutionProviderFactory_DML(device_id));
#endif
} else if (type == kNnapiExecutionProvider) {
#if defined(USE_NNAPI)