import triton files when triton is supported and installed

deepspeedai · Feb 1, 2025 · ec8e8e1 · ec8e8e1
1 parent 241bffd
commit ec8e8e1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/model_implementations/transformers/ds_transformer.py b/deepspeed/model_implementations/transformers/ds_transformer.py
@@ -14,7 +14,7 @@
 from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp
 from deepspeed.accelerator import get_accelerator
 import deepspeed
-if deepspeed.HAS_TRITON:
+if deepspeed.HAS_TRITON and get_accelerator().is_triton_supported():
     from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP
     from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention