def to_device(data, device):
    """ Transfer data to device.

    Parameters
    ----------
    data: tensor or list of tensor
        the data to be transfered.
    device: torch.device
        the device to work on.

    Returns
    -------
    out: tensor or list of tensor
        the transfered data.
    """
    if isinstance(data, list):
        return [tensor.to(device) for tensor in data]
    else:
        return data.to(device)