Training Keras model on Google Colab TPU: UnimplementedError: Compilation failure: HLO get-dimension-size is not implemented on TPU

Question

I\'m trying to train a Keras model on Google Colab with TPU and getting the error mentioned in the title.

The relevant code is as follows:

tpu = tf.dis