def center_crop(batch, size):
    """
    Crop the center of a 4D images tensor

    Args:
        batch (4D images tensor): the tensor to crop
        size ((int, int)): size of the resulting image as (height, width)

    Returns:
        The cropped image
    """
    if batch.shape[-2] == size[0] and batch.shape[-1] == size[1]:
        return batch
    y_off = (batch.shape[-2] - size[0]) // 2
    x_off = (batch.shape[-1] - size[1]) // 2
    return batch[..., y_off:y_off + size[0], x_off:x_off + size[1]]