Mike0021/zonos2

1from __future__ import annotations2 3from functools import lru_cache4from typing import Tuple5 6 7@lru_cache(maxsize=None)8def _get_torch_cuda_version() -> Tuple[int, int] | None:9    import torch10    import torch.version11 12    if not torch.cuda.is_available() or not torch.version.cuda:13        return None14    return torch.cuda.get_device_capability()15 16 17def is_arch_supported(major: int, minor: int = 0) -> bool:18    arch = _get_torch_cuda_version()19    if arch is None:20        return False21    return arch >= (major, minor)22 23 24def is_sm90_supported() -> bool:25    return is_arch_supported(9, 0)26 27 28def is_sm100_supported() -> bool:29    return is_arch_supported(10, 0)30