"OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed in Graph execution." error on using MirroredStrategy

I am trying to train a model on multiple GPUs using mirroredStrategy. I have been following the below link: https://www.tensorflow.org/tutorials/distribute/custom_training

Now, I have followed the process almost exactly. But still I am getting the following error:

INFO:tensorflow:Error reported to Coordinator: iterating over `tf.Tensor` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.

I have created the dataset like this:

dataset = tf.data.Dataset.from_tensor_slices(list_ds)
train_dataset = dataset.batch(GLOBAL_BATCH_SIZE, drop_remainder=True)
train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)

list_ds is a list of lists of strings.

And this is how the my training loop looks:

@tf.function
def distributed_train_step(dataset_inputs):
    strategy.experimental_run_v2(train_step, args=dataset_inputs)

for epoch in range(EPOCHS):       
    for batch in train_dist_dataset:
        distributed_train_step(batch)

I also ran the exact same code provided in the link and it ran successfully but when I try the same on my model, I get the error.