TPU fails with BERT for semantic similarity in Keras but works on GPU

Question

I am experimenting with BERT for semantic similarity and I follow this code here https://keras.io/examples/nlp/semantic_similarity_with_bert/

Everything runs smoothly