diff --git a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py index e8d30b5ff..5e487141d 100644 --- a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py +++ b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py @@ -4,11 +4,8 @@ def get_backend() -> str: """Detect the backend by inspecting torch.""" import torch - if hasattr(torch.backends, "tpu"): - # torch_tpu sets torch.backends.tpu when it is imported (via - # torch's device-backend autoload), regardless of whether TPU - # hardware is present — analogous to torch.version.cuda being - # set on CUDA builds without a GPU. + acc = torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None + if acc is not None and acc.type == "tpu": return "tpu" elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron diff --git a/kernels/src/kernels/backends.py b/kernels/src/kernels/backends.py index cd4473421..b8e20a8a6 100644 --- a/kernels/src/kernels/backends.py +++ b/kernels/src/kernels/backends.py @@ -217,12 +217,8 @@ def _backend() -> Backend: if has_torch: import torch - if hasattr(torch.backends, "tpu"): - # torch_tpu sets torch.backends.tpu when it is imported (via - # torch's device-backend autoload), regardless of whether TPU - # hardware is present — analogous to torch.version.cuda being - # set on CUDA builds without a GPU. The hardware-gated - # torch.tpu namespace only appears on hosts with TPU devices. + acc = torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None + if acc is not None and acc.type == "tpu": return TPU() elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron