Compressor/backend/LabelingPreTraining/tspulse/device.py

"""Device selection for TSPulse training/inference.

Supports CUDA, Apple MPS, Huawei Ascend NPU and CPU, chosen automatically in
that priority order. Set TSPULSE_DEVICE to force a specific backend.
"""

import os

import torch


def get_device() -> torch.device:
    forced = os.getenv("TSPULSE_DEVICE", "").strip().lower()
    if forced:
        return torch.device(forced)

    if torch.cuda.is_available():
        return torch.device("cuda")

    if torch.backends.mps.is_available():
        return torch.device("mps")

    try:
        import torch_npu  # noqa: F401

        if torch.npu.is_available():
            return torch.device("npu")
    except ImportError:
        pass

    return torch.device("cpu")


def describe(device: torch.device) -> str:
    if device.type == "cuda":
        name = torch.cuda.get_device_name(device)
        return f"CUDA ({name})"
    if device.type == "mps":
        return "Apple MPS"
    if device.type == "npu":
        return "Huawei Ascend NPU"
    return "CPU"