def only_first_4_patches(tensor, n_imgs=4):
    """
    Batch truncation function
    :param tensor: batch
    :param n_imgs: number of patches to keep
    :return: truncated batch
    """
    return tensor[0:n_imgs]