condition [32,1,1,1], then [16,64,64,3], and else [16,64,64,3] must be broadcastable

I have issue in code. I run code after 1 epochs program terminate and also change batch size but can't work. I also search on google this error but can't find any solution.Any Expert Please Checkout my error

Full Code Link Colab Link :https://colab.research.google.com/drive/1kPDPIXliTR1Wxc8vzJ5wjgxoM7fzgnwG?usp=sharing

self.probability = tf.Variable(0.0)

self.augmenter = keras.Sequential(
    [
        layers.InputLayer(input_shape=(image_size, image_size, 3)),
        # blitting/x-flip:
        layers.RandomFlip("horizontal"),
        # blitting/integer translation:
        layers.RandomTranslation(
            height_factor=max_translation,
            width_factor=max_translation,
            interpolation="nearest",
        ),
        # geometric/rotation:
        layers.RandomRotation(factor=max_rotation),
        # geometric/isotropic and anisotropic scaling:
        layers.RandomZoom(
            height_factor=(-max_zoom, 0.0), width_factor=(-max_zoom, 0.0)
        ),
    ],
    name="adaptive_augmenter",
)

def call(self, images, training): if training: augmented_images = self.augmenter(images, training)

    augmentation_values = tf.random.uniform(
        shape=(batch_size, 1, 1, 1), minval=0.0, maxval=1.0
    )
    augmentation_bools = tf.math.less(augmentation_values, self.probability)

    images = tf.where(augmentation_bools, augmented_images, images)
return images

def update(self, real_logits): current_accuracy = tf.reduce_mean(step(real_logits))

accuracy_error = current_accuracy - target_accuracy
self.probability.assign(
    tf.clip_by_value(
        self.probability + accuracy_error / integration_steps, 0.0, 1.0
    )
)

Error

InvalidArgumentError: condition [8,1,1,1], then [32,128,128,3], and else [32,128,128,3] must be broadcastable [[node adaptive_augmenter/SelectV2 (defined at :46) ]] [Op:__inference_train_function_5713]

Errors may have originated from an input operation. Input Source operations connected to node adaptive_augmenter/SelectV2: adaptive_augmenter/Less (defined at :44) IteratorGetNext (defined at :1) adaptive_augmenter/adaptive_augmenter/random_zoom/transform/ImageProjectiveTransformV3 (defined at :38)

Function call stack: train_function