If the input is a tensor, return it. Otherwise, return the first element.
Input batch.
A tensor, used as the input to a model.