Colab TPU: InvalidArgumentError: Unable to parse tensor proto [Op:ComputeBatchSize]

Question

I'm trying to use TPU on Colab and I got this error when calling model.fit()

The error message is simply: InvalidArgumentError: Unable to parse tensor proto [Op:ComputeBatchSize]

Since it has mentioned BatchSize, I tried to set different batch size but didn't work. Here's my data pipeline:

gcs_pattern = 'gs://diem-bucket/*.tfrecords' 
filenames = tf.io.gfile.glob(gcs_pattern) ds_train = tf.data.Dataset.list_files(filenames, shuffle= True)

ds_train = ds_train.shuffle(
                    buffer_size = len(tfrcfiles), 
                    reshuffle_each_iteration = True,
                ).interleave(
                    map_func = lambda x: tf.data.TFRecordDataset(x),
                    cycle_length = len(tfrcfiles),
                    num_parallel_calls = tf.data.AUTOTUNE,
                    deterministic=False
                ).batch(
                    8
                ).prefetch(
                    8
                ).map(
                    batch_parse_func, 
                    num_parallel_calls = tf.data.AUTOTUNE
                )

I could barely found any information about this error. Can someone share some advice?