import torch


def select_device(cuda_number=0):
    """
    Returns a device for the model
    :param cuda_number: Id of the cuda, if available.
    :return: a device for the model
    """
    cuda_name = "cuda:"+str(cuda_number)
    device = torch.device(cuda_name if torch.cuda.is_available() else "cpu")
    return device