(feat): adding intel gpus support

This commit is contained in:
rahulunair
2024-05-13 21:21:05 +00:00
parent 1b5130694a
commit 6be88fa81b
3 changed files with 89 additions and 0 deletions

View File

@@ -27,6 +27,12 @@ def detect_device():
"""
Detects the appropriate device to run on, and return the device and dtype.
"""
try:
import intel_extension_for_pytorch
if torch.xpu.is_available():
return torch.device("xpu"), torch.float32
except ImportError:
pass
if torch.cuda.is_available():
return torch.device("cuda"), torch.float16
elif torch.backends.mps.is_available():