Merge pull request #146 from nateraw/onnx-gpu

use cuda with onnx if available
spotify · Nov 4, 2024 · 53f6738 · 53f6738
2 parents 9991303 + b7e349d
commit 53f6738
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"):