Explicitly specify cuda device for onnxruntime

BreezeWhite · Nov 16, 2024 · 76894b4 · 76894b4
1 parent dd00ca1
commit 76894b4
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/oemer/inference.py b/oemer/inference.py
@@ -46,13 +46,17 @@ def inference(
         output_shape = model.output_shape
     else:
         import onnxruntime as rt
+        import torch
 
         onnx_path = os.path.join(model_path, "model.onnx")
         metadata = pickle.load(open(os.path.join(model_path, "metadata.pkl"), "rb"))
         if sys.platform == "darwin":
             providers = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
         else:
-            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            providers = [
+                ("CUDAExecutionProvider", {"device_id": 0}),
+                "CPUExecutionProvider",
+            ]
         sess = rt.InferenceSession(onnx_path, providers=providers)
         output_names = metadata["output_names"]
         input_shape = metadata["input_shape"]