From e1ed0fde2bbcf2b56898a0ca456f2fc2981c4f64 Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Tue, 8 Sep 2020 11:13:50 -0700 Subject: [PATCH] Prevent registering both DML and CUDA EPs in an ML op test (#5078) --- onnxruntime/python/onnxruntime_pybind_state.cc | 1 + .../test/python/onnxruntime_test_python_mlops.py | 12 +++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 723debbe24..fea461ed98 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -571,6 +571,7 @@ void RegisterExecutionProviders(InferenceSession* sess, const std::vectorGetSessionOptions().enable_cpu_mem_arena)); #endif + } else if (type == kDmlExecutionProvider) { } else { // unknown provider throw std::runtime_error("Unknown Provider Type: " + type); diff --git a/onnxruntime/test/python/onnxruntime_test_python_mlops.py b/onnxruntime/test/python/onnxruntime_test_python_mlops.py index 0728c825ab..f30675a702 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_mlops.py +++ b/onnxruntime/test/python/onnxruntime_test_python_mlops.py @@ -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)