[Bugfix] Improves compatibility when checking for MPS availability in different PyTorch builds. (tile-ai#1051)

Author: LeiWang1999

tilelang/utils/device.py

@@ -1,7 +1,14 @@
 import torch
 
 IS_CUDA = torch.cuda.is_available()
-IS_MPS = torch.mps.is_available()
+
+IS_MPS = False
+try:
+    IS_MPS = torch.backends.mps.is_available()
+except AttributeError:
+    print("MPS backend is not available in this PyTorch build.")
+except Exception as e:
+    print(f"An unexpected error occurred while checking MPS availability: {e}")
 
 
 def get_current_device():