fix triton version check (#3345)

* fiix triton version check * add xpu check
huggingface · Jan 17, 2025 · 02d2561 · 02d2561
1 parent fbfa53b
commit 02d2561
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
@@ -229,6 +229,8 @@ def is_timm_available():
 
 
 def is_triton_available():
+    if is_xpu_available():
+        return _is_package_available("triton", "pytorch-triton-xpu")
     return _is_package_available("triton")