WARNING:tensorflow:AutoGraph could not transform <function _preprocess at 0x7f8ff80cd160> and will run it as-is

I'm trying to run the code found in https://www.tensorflow.org/probability/examples/Probabilistic_Layers_VAE .

I'm using Python version 3.9 and my TensorFlow version is >2.0. The code is as follows:

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()

import tensorflow_datasets as tfds
import tensorflow_probability as tfp


tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions

datasets, datasets_info = tfds.load(name='mnist', with_info=True, as_supervised=False)

def _preprocess(sample):
    image = tf.cast(sample['image'], tf.float32) / 255 #Scale to [0, 1]
    image = image < tf.random.uniform(tf.shape(image)) #Gives 0, 1 when compared to a random number
    return image, image

train_dataset = (datasets['train']
                .map(_preprocess)
                .batch(256)
                .prefetch(tf.data.experimental.AUTOTUNE)
                .shuffle(int(10e3)))

What I get is the following warning:

WARNING:tensorflow:AutoGraph could not transform <function _preprocess at 0x7f8ff80cd160> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: module 'gast' has no attribute 'Index'
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert

The warning is related to the last part of the code but I can't tell if it's going to potentially affect how the code runs. If it won't affect it, is there a way to consistently remove such warnings?