Fix ORT execution provider test for onnxruntime v1.16.0 (#1406)

huggingface · Sep 21, 2023 · 6d1ae0e · 6d1ae0e
1 parent 3c4ad78
commit 6d1ae0e
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py
@@ -357,25 +357,22 @@ def test_missing_execution_provider(self):
 
         if not is_onnxruntime_gpu_installed:
             for provider in ["CUDAExecutionProvider", "TensorrtExecutionProvider"]:
-                with self.assertRaises(ImportError) as cm:
+                with self.assertRaises(ValueError) as cm:
                     _ = ORTModel.from_pretrained(self.ONNX_MODEL_ID, provider=provider)
 
-                self.assertTrue(
-                    f"Asked to use {provider}, but `onnxruntime-gpu` package was not found." in str(cm.exception)
-                )
+            self.assertTrue("but the available execution providers" in str(cm.exception))
+
         else:
             logger.info("Skipping CUDAExecutionProvider/TensorrtExecutionProvider without `onnxruntime-gpu` test")
 
         # need to install first onnxruntime-gpu, then onnxruntime for this test to pass,
         # thus overwritting onnxruntime/capi/_ld_preload.py
         if is_onnxruntime_installed and is_onnxruntime_gpu_installed:
             for provider in ["CUDAExecutionProvider", "TensorrtExecutionProvider"]:
-                with self.assertRaises(ImportError) as cm:
+                with self.assertRaises(ValueError) as cm:
                     _ = ORTModel.from_pretrained(self.ONNX_MODEL_ID, provider=provider)
 
-                self.assertTrue(
-                    "`onnxruntime-gpu` is installed, but GPU dependencies are not loaded." in str(cm.exception)
-                )
+                self.assertTrue("but the available execution providers" in str(cm.exception))
         else:
             logger.info("Skipping double onnxruntime + onnxruntime-gpu install test")