diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 74312389c9..c111cce6b9 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -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()