How can I get different behaviour within Tensorflow call function during training?

I've got a preprocessing layer pp_layer that is called during training. I want to find a way to change the behaviour of my subclassed model such that it also calls pp_layer during inference, but won't call it twice during training.

Currently my code looks like this:

@tf.function
    def train_step(self, data):
        x, y = data
        x = self.pp_layer(x)
        y = tf.cast(y, tf.int32)
...
@tf.function(input_signature=[tf.TensorSpec([None, 256, 256, 3], tf.float32)])
    def call(self, data):
        data = self.pp_layer(data)

How can I modify call so that it is aware of when training is taking place?

I've tried adding a training=False argument but this causes problems as I have my @tf.function(input_signature=[tf.TensorSpec([None, 256, 256, 3], tf.float32)]) signature above call to set the input shapes for exporting the model, so Tensorflow complains that the input arguments no longer match the signature.