Add cuda support when loading local onnx model

onnx · Dec 13, 2024 · 004cb07 · 004cb07
1 parent 63f11db
commit 004cb07
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/turnkeyml/llm/tools/ort_genai/oga.py b/src/turnkeyml/llm/tools/ort_genai/oga.py
@@ -35,7 +35,7 @@
 oga_model_builder_cache_path = "model_builder"
 
 # Mapping from processor to executiion provider, used in pathnames and by model_builder
-execution_providers = {"cpu": "cpu", "npu": "npu", "igpu": "dml"}
+execution_providers = {"cpu": "cpu", "npu": "npu", "igpu": "dml", "cuda": "cuda"}
 
 
 class OrtGenaiTokenizer(TokenizerAdapter):
@@ -248,7 +248,7 @@ def parser(add_help: bool = True) -> argparse.ArgumentParser:
         parser.add_argument(
             "-d",
             "--device",
-            choices=["igpu", "npu", "cpu"],
+            choices=["igpu", "npu", "cpu", "cuda"],
             default="igpu",
             help="Which device to load the model on to (default: igpu)",
         )