Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021"""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 amachine 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")