mighty.utils.common.input_from_batch

mighty.utils.common.input_from_batch(batch)

If the input is a tensor, return it. Otherwise, return the first element.

Parameters:

batchtorch.Tensor or tuple of torch.Tensor

Input batch.

Returns:

torch.Tensor

A tensor, used as the input to a model.