Function is_torch_tpu_available

src/transformers/utils/import_utils.py:526–542

(check_device: bool = False)

@lru_cache
def is_torch_tpu_available(check_device: bool = False) -> bool:
    import torch

    if importlib.util.find_spec("torch_tpu") is None:
        return False

    if check_device:
        try:
            import torch_tpu  # noqa: F401

            if hasattr(torch, "tpu") and torch.tpu.is_available():
                return torch.tpu.device_count() >= 1
            return False
        except RuntimeError:
            return False

    return hasattr(torch, "tpu") and torch.tpu.is_available()


@lru_cache