Commit 507a4afc authored by Matthew Douglas's avatar Matthew Douglas
Browse files

Improve guard for triton compatibility (#1497)

parent 6de46d99
import importlib import functools
@functools.lru_cache(None)
def is_triton_available(): def is_triton_available():
return importlib.util.find_spec("triton") is not None try:
# torch>=2.2.0
from torch.utils._triton import has_triton, has_triton_package
return has_triton_package() and has_triton()
except ImportError:
from torch._inductor.utils import has_triton
return has_triton()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment