src/utils/device.py

"""
Device auto-detection utility.

Selects the best available compute device in priority order:
CUDA (NVIDIA GPU) > MPS (Apple Silicon) > CPU.
Provides a single entry point so model and tensor placement
is consistent across the codebase.
"""

import torch

def get_default_device(force_cpu: bool = False):
    """
    Automatically selects the best available device.
    Priority: CUDA > MPS (Apple Silicon) > CPU
    
    Args:
        force_cpu: If True, forces CPU usage (useful for debugging).
    """
    if force_cpu:
        return torch.device("cpu")

    if torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available():
        return torch.device("mps")
    else:
        return torch.device("cpu")

def move_to_device(model, device=None):
    """
    Helper to move model to device and print status.
    """
    if device is None:
        device = get_default_device()
        
    model.to(device)
    print(f"🚀 Model moved to: {device}")