import torch


def get_pytorch_device() -> torch.device:
    """Checks if a CUDA enabled GPU is available, and returns the
    approriate device, either CPU or GPU.

    Returns
    -------
    device : torch.device
    """

    device = torch.device("cpu")
    if torch.cuda.is_available():
        device = torch.device("cuda:0")

    return device