Skip to content

Commit

Permalink
Update is_mlu_available (#1537)
Browse files Browse the repository at this point in the history
* Update is_mlu_available

to adapt torch_mlu main, the torch.is_mlu_available method is removed

* Update utils.py

* Update utils.py
  • Loading branch information
ClowDragon authored May 30, 2024
1 parent d1f1aab commit 85c83ba
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion mmengine/device/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@
except Exception:
IS_NPU_AVAILABLE = False

try:
import torch_mlu # noqa: F401
IS_MLU_AVAILABLE = hasattr(torch, 'mlu') and torch.mlu.is_available()
except Exception:
IS_MLU_AVAILABLE = False

try:
import torch_dipu # noqa: F401
IS_DIPU_AVAILABLE = True
Expand Down Expand Up @@ -64,7 +70,7 @@ def is_npu_available() -> bool:

def is_mlu_available() -> bool:
"""Returns True if Cambricon PyTorch and mlu devices exist."""
return hasattr(torch, 'is_mlu_available') and torch.is_mlu_available()
return IS_MLU_AVAILABLE


def is_mps_available() -> bool:
Expand Down

0 comments on commit 85c83ba

Please sign in to comment.