fix: report correct device count for intel xpu
This commit is contained in:
parent
625a0e32c4
commit
dc36fcc3de
|
@ -141,6 +141,8 @@ def get_device_count() -> int:
|
|||
return torch.npu.device_count()
|
||||
elif is_torch_cuda_available():
|
||||
return torch.cuda.device_count()
|
||||
elif is_torch_xpu_available():
|
||||
return torch.xpu.device_count()
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
|
Loading…
Reference in New Issue