Only allow CUDAExecutionProvider if available

wkentaro · Feb 12, 2024 · e0af059 · e0af059
1 parent e896da9
commit e0af059
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/osam/_models/_base.py b/osam/_models/_base.py
@@ -62,7 +62,10 @@ def __init__(self):
             try:
                 # Try to use all of the available providers e.g., cuda, tensorrt.
                 if providers is None:
-                    providers = onnxruntime.get_available_providers()
+                    if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
+                        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+                    else:
+                        providers = ["CPUExecutionProvider"]
                 # Suppress all the error messages from the missing providers.
                 with _contextlib.suppress():
                     inference_session = onnxruntime.InferenceSession(