tiny fix

2024-08-20 00:10:52 +08:00
parent 15be296347
commit c8b4c7fee5
2 changed files with 4 additions and 4 deletions
--- a/src/llamafactory/extras/misc.py
+++ b/src/llamafactory/extras/misc.py
@@ -137,12 +137,12 @@ def get_device_count() -> int:
    r"""
    Gets the number of available GPU or NPU devices.
    """
-    if is_torch_npu_available():
+    if is_torch_xpu_available():
+        return torch.xpu.device_count()
+    elif is_torch_npu_available():
        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