import torch


def _set_device(gpu="0"):

    """
    Parameters:
    -----------
    gpu
        the number (typically 0) indicating the GPU attached to the running instance.
        
    Returns:
    --------
    device
        object to which torch.Tensors can be sent to for calculations.
    
    """

    device = torch.device("cuda:" + gpu if torch.cuda.is_available() else "cpu")

    return device