import torch


def input_from_batch(batch):
    """
    If the input is a tensor, return it. Otherwise, return the first element.

    Parameters
    ----------
    batch : torch.Tensor or tuple of torch.Tensor
        Input batch.

    Returns
    -------
    torch.Tensor
        A tensor, used as the input to a model.

    """
    if isinstance(batch, torch.Tensor):
        # unsupervised learning, no labels
        return batch
    # iterable
    return batch[0]