import torch
import logging


def get_device(device: str = "CPU") -> torch.device:
    """
    Get torch device.

    Args:
        device: Specifies device to use

    Returns:
        device: Torch device if found

    Raises:
        ValueError: If GPU was selected, but not available
    """
    if device == "GPU":
        if torch.cuda.is_available():
            torch_device = torch.device("cuda")
        else:
            raise RuntimeError("GPU was selected, but no CUDA GPUs are available")
    else:
        torch_device = torch.device("cpu")

    logging.info(f"Running on {device}")
    return torch_device