import torch


def set_device(
    var, 
    use_cuda: bool
):
    """
    Gets tensor or nn variable ready for computation on a device.

    params:
        <tensor or nn> var: Tensor or model
        <bool> use_cuda: Whether to use GPU for computation or not

    returns:
        <tensor or nn> var: The same tensor or model, prepared for computation on the selected device
    """
    if use_cuda == True:
        if torch.cuda.is_available():
            var = var.cuda()
        else:
            var = var.cpu()
    else:
        var = var.cpu()

    return var