Prevent registering both DML and CUDA EPs in an ML op test (#5078)

This commit is contained in:
Hariharan Seshadri 2020-09-08 11:13:50 -07:00 committed by GitHub
parent 8d91d4ff36
commit e1ed0fde2b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 1 deletions

View file

@ -571,6 +571,7 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vector<std::s
RegisterExecutionProvider(
sess, *onnxruntime::CreateExecutionProviderFactory_ArmNN(sess->GetSessionOptions().enable_cpu_mem_arena));
#endif
} else if (type == kDmlExecutionProvider) {
} else {
// unknown provider
throw std::runtime_error("Unknown Provider Type: " + type);

View file

@ -133,7 +133,17 @@ class TestInferenceSession(unittest.TestCase):
np.testing.assert_allclose(output_expected, res[0], rtol=1e-05, atol=1e-08)
def test_run_model_mlnet(self):
sess = onnxrt.InferenceSession(get_name("mlnet_encoder.onnx"))
available_providers = onnxrt.get_available_providers()
# The Windows GPU CI pipeline builds the wheel with both CUDA and DML enabled and ORT does not support cases
# where one node is asigned to CUDA and one node to DML as it doesn't have the data transfer capabilities to deal with
# potentially different device memory. Hence, use a session with only DML and CPU (excluding CUDA) for this test as it breaks
# with both CUDA and DML registered.
if ('CUDAExecutionProvider' in available_providers and 'DmlExecutionProvider' in available_providers):
sess = onnxrt.InferenceSession(get_name("mlnet_encoder.onnx"), None, ['DmlExecutionProvider', 'CPUExecutionProvider'])
else:
sess = onnxrt.InferenceSession(get_name("mlnet_encoder.onnx"))
names = [_.name for _ in sess.get_outputs()]
self.assertEqual(['C00', 'C12'], names)
c0 = np.array([5.], dtype=np.float32).reshape(1, 1)