Add support for onnxruntime-directml in import utilities

huggingface · Feb 2, 2025 · 66ed0ff · 66ed0ff
1 parent e6d15fe
commit 66ed0ff
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum/utils/import_utils.py b/optimum/utils/import_utils.py
@@ -113,6 +113,8 @@ def _is_package_available(
         "onnxruntime-migraphx",
         "ort-migraphx-nightly",
         "ort-rocm-nightly",
+        # For DirectML
+        "onnxruntime-directml",
     ],
 )
 _tf_available, _tf_version = _is_package_available(