mini-modern-llm / src /device.py
hey-shiv's picture
Upload folder using huggingface_hub
43d2e8b verified
"""Device selection helpers."""
from __future__ import annotations
import torch
def get_default_device() -> torch.device:
"""Prefer CUDA, then Apple MPS, then CPU."""
if torch.cuda.is_available():
return torch.device("cuda")
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")