import torch


def check_cuda(requestCuda=True):
    """
    Check if CUDA is available, set and return the device

    Inputs:
        USE_CUDA: use Cuda if True and Cuda is available, otherwise use cpu
    
    Returns:
        the device selected ('cpu' or the cuda device)
    """
    print('Requested CUDA:',requestCuda)
    if requestCuda==True:
        print('There are {} CUDA Devices installed on this system:'.format(torch.cuda.device_count()))
        use_cuda = torch.cuda.is_available()
        device = torch.device("cuda:0" if use_cuda else "cpu")
        print('CUDA Device:',device,torch.cuda.get_device_name(device))
        print('Use CUDA? ',use_cuda)
    else:
        use_cuda = False
        device = torch.device('cpu')
        print('Using CPU as torch device')    
    return device