import torch


def get_device(request_gpu):
    """
    Return torch.device depending on request_gpu and GPU availability

    :param request_gpu: Set to True for requesting use of GPU
    :return torch.device:
    """
    if request_gpu:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if device == torch.device("cpu"):
            print("WARN: GPU requested, but is not available. Use CPU instead.")
    else:
        device = torch.device("cpu")

    return device