We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
1 parent d1f1aab commit 85c83baCopy full SHA for 85c83ba
mmengine/device/utils.py
@@ -16,6 +16,12 @@
16
except Exception:
17
IS_NPU_AVAILABLE = False
18
19
+try:
20
+ import torch_mlu # noqa: F401
21
+ IS_MLU_AVAILABLE = hasattr(torch, 'mlu') and torch.mlu.is_available()
22
+except Exception:
23
+ IS_MLU_AVAILABLE = False
24
+
25
try:
26
import torch_dipu # noqa: F401
27
IS_DIPU_AVAILABLE = True
@@ -64,7 +70,7 @@ def is_npu_available() -> bool:
64
70
65
71
def is_mlu_available() -> bool:
66
72
"""Returns True if Cambricon PyTorch and mlu devices exist."""
67
- return hasattr(torch, 'is_mlu_available') and torch.is_mlu_available()
73
+ return IS_MLU_AVAILABLE
68
74
69
75
76
def is_mps_available() -> bool:
0 commit comments