Tensorflow: Custom data augmentation

Question

I'm trying to define a custom data augmentation layer. My goal is to call the existing tf.keras.layers.RandomZoom, with a probability.

This is what I did:

class random_zoom_layer(tf.keras.layers.Layer):
    def __init__(self, probability=0.5, **kwargs):
        super().__init__(**kwargs)
        self.probability = probability

    def call(self, x):
        if tf.random.uniform([]) < self.probability:
            return tf.keras.layers.RandomZoom(height_factor=(-0.1, 0.1), width_factor=(-0.1, 0.1), fill_mode='constant')(x)
        else:
            return x


data_augmentation = tf.keras.Sequential([
    tf.keras.layers.experimental.preprocessing.Normalization(),
    random_zoom_layer(probability=0.2)
])

But during training, I receive this error:

tensorflow.python.framework.errors_impl.NotFoundError: 2 root error(s) found.
  (0) NOT_FOUND:  2 root error(s) found.
  (0) NOT_FOUND:  Resource localhost/_AnonymousVar10/class tensorflow::Var does not exist.
     [[{{node sequential_1/random_zoom_layer/cond/random_zoom/stateful_uniform/RngReadAndSkip}}]]
     [[sequential_1/random_zoom_layer/cond/then/_0/sequential_1/random_zoom_layer/cond/random_zoom/stateful_uniform_1/RngReadAndSkip/_15]]
  (1) NOT_FOUND:  Resource localhost/_AnonymousVar10/class tensorflow::Var does not exist.
     [[{{node sequential_1/random_zoom_layer/cond/random_zoom/stateful_uniform/RngReadAndSkip}}]]

I would really appreciate some help!