Skip to content

Commit 1f4ffdb

Browse files
authored
[Bugfix] Improves compatibility when checking for MPS availability in different PyTorch builds. (#1051)
1 parent e3742d3 commit 1f4ffdb

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

tilelang/utils/device.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,14 @@
11
import torch
22

33
IS_CUDA = torch.cuda.is_available()
4-
IS_MPS = torch.mps.is_available()
4+
5+
IS_MPS = False
6+
try:
7+
IS_MPS = torch.backends.mps.is_available()
8+
except AttributeError:
9+
print("MPS backend is not available in this PyTorch build.")
10+
except Exception as e:
11+
print(f"An unexpected error occurred while checking MPS availability: {e}")
512

613

714
def get_current_device():

0 commit comments

Comments
 (0)