Update __init__.py

This commit is contained in:
Daniel Han 2025-10-14 05:54:32 -07:00
commit 10a5b76623

View file

@ -82,7 +82,18 @@ def get_device_type():
return "cuda"
elif hasattr(torch, "xpu") and torch.xpu.is_available():
return "xpu"
raise NotImplementedError("Unsloth currently only works on NVIDIA GPUs and Intel GPUs.")
# Check torch.accelerator
if hasattr(torch, "accelerator"):
if not torch.accelerator.is_available():
raise NotImplementedError("Unsloth cannot find any torch accelerators? You need a GPU.")
accelerator = str(torch.accelerator.current_accelerator())
if accelerator in ("cuda", "xpu", "hip"):
raise RuntimeError(
f"Unsloth: Weirdly `torch.cuda.is_available()`, `torch.xpu.is_available()` and `is_hip` all failed.\n"\
f"But `torch.accelerator.current_accelerator()` works with it being = `{accelerator}`\n"\
f"Please reinstall torch - it's most likely broken :("
)
raise NotImplementedError("Unsloth currently only works on NVIDIA, AMD and Intel GPUs.")
pass
DEVICE_TYPE : str = get_device_type()