Update __init__.py
This commit is contained in:
parent
f473b69635
commit
10a5b76623
1 changed files with 12 additions and 1 deletions
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue