Skip to content

Commit 85c83ba

Browse files
authored
Update is_mlu_available (#1537)
* Update is_mlu_available to adapt torch_mlu main, the torch.is_mlu_available method is removed * Update utils.py * Update utils.py
1 parent d1f1aab commit 85c83ba

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

mmengine/device/utils.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,12 @@
1616
except Exception:
1717
IS_NPU_AVAILABLE = False
1818

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+
1925
try:
2026
import torch_dipu # noqa: F401
2127
IS_DIPU_AVAILABLE = True
@@ -64,7 +70,7 @@ def is_npu_available() -> bool:
6470

6571
def is_mlu_available() -> bool:
6672
"""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
6874

6975

7076
def is_mps_available() -> bool:

0 commit comments

Comments
 (0)