Skip to content

Commit 507a4af

Browse files
Improve guard for triton compatibility (#1497)
1 parent 6de46d9 commit 507a4af

File tree

1 file changed

+11
-2
lines changed

1 file changed

+11
-2
lines changed

bitsandbytes/triton/triton_utils.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,14 @@
1-
import importlib
1+
import functools
22

33

4+
@functools.lru_cache(None)
45
def is_triton_available():
5-
return importlib.util.find_spec("triton") is not None
6+
try:
7+
# torch>=2.2.0
8+
from torch.utils._triton import has_triton, has_triton_package
9+
10+
return has_triton_package() and has_triton()
11+
except ImportError:
12+
from torch._inductor.utils import has_triton
13+
14+
return has_triton()

0 commit comments

Comments
 (0)