"""Device selection for training/eval: cuda (includes ROCm), then mps, then cpu.

Override with DSEM_DEVICE=cpu|mps|cuda|cuda:N when needed (e.g. forcing cpu on a
machine with a broken GPU stack, or pinning a specific GPU without CUDA_VISIBLE_DEVICES).
"""

import os

import torch


def get_device() -> torch.device:
    override = os.environ.get("DSEM_DEVICE")
    if override:
        return torch.device(override)
    if torch.cuda.is_available():  # true for both CUDA and ROCm builds
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")