use cuda with onnx if available

spotify · Oct 23, 2024 · b7e349d · b7e349d
1 parent 9991303
commit b7e349d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/basic_pitch/inference.py b/basic_pitch/inference.py
@@ -129,7 +129,10 @@ def __init__(self, model_path: Union[pathlib.Path, str]):
             present.append("ONNX")
             try:
                 self.model_type = Model.MODEL_TYPES.ONNX
-                self.model = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
+                providers = ["CPUExecutionProvider"]
+                if "CUDAExecutionProvider" in ort.get_available_providers():
+                    providers.insert(0, "CUDAExecutionProvider")
+                self.model = ort.InferenceSession(str(model_path), providers=providers)
                 return
             except Exception as e:
                 if str(model_path).endswith(".onnx"):