From 8618922a33b6e83118c2a6fd53c7502c76e7b8ca Mon Sep 17 00:00:00 2001
From: a-One-Fan <espamholding@gmail.com>
Date: Thu, 4 May 2023 21:14:07 +0300
Subject: [PATCH] Implement correct XPU device count Forgot to do that
---
tortoise/utils/device.py | 3 +++
1 file changed, 3 insertions(+)
diff --git a/tortoise/utils/device.py b/tortoise/utils/device.py
index eddfcea..b4d51ea 100755
--- a/tortoise/utils/device.py
+++ b/tortoise/utils/device.py
@@ -128,6 +128,9 @@ def get_device_count(name=get_device_name()):
if name == "dml":
import torch_directml
return torch_directml.device_count()
+ if name == "xpu":
+ import intel_extension_for_pytorch
+ return torch.xpu.device_count()
return 1