src/utils/device_utils.py

import torch

def get_device():
    """
    Determine the available device for computation.
    Returns either CUDA device if available, or CPU.
    """
    if torch.cuda.is_available():
        return torch.device('cuda')
    return torch.device('cpu')