import torch


def reshape_batch(batch, column='examples'):
    """ This reshapes the batch to have an extra dimension corresponding to the input channels so we can apply the torch.nn.Conv2d operation in our model. """
    shape = batch[column].shape
    new_shape = torch.Size([shape[0], 1, shape[1], shape[2]])
    batch[column] = batch[column].reshape(new_shape)
    return batch