[NPU] Bug fix in device detect (#14137)

This commit is contained in:
Mayyyy
2025-12-25 00:18:39 +08:00
committed by GitHub
parent 8bf7f240b6
commit 2fb3160598

View File

@@ -144,7 +144,15 @@ def is_xpu() -> bool:
@lru_cache(maxsize=1)
def is_npu() -> bool:
return hasattr(torch, "npu") and torch.npu.is_available()
if not hasattr(torch, "npu"):
return False
if not torch.npu.is_available():
raise RuntimeError(
"torch_npu detected, but NPU device is not available or visible."
)
return True
@lru_cache(maxsize=1)
@@ -1854,7 +1862,7 @@ def get_device(device_id: Optional[int] = None) -> str:
return "xpu"
return "xpu:{}".format(device_id)
if hasattr(torch, "npu") and torch.npu.is_available():
if is_npu():
if device_id == None:
return "npu"
return "npu:{}".format(device_id)