Error while compiling on TPU

[gaphex]

I am using tf 1.13 on Google Colab.
When compiling the model for TPU I got an issue:

training_model_tpu = tf.contrib.tpu.keras_to_tpu_model(
    training_model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    164     cls = module_objects.get(class_name)
    165     if cls is None:
--> 166       raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    167   return (cls, config['config'])
    168 

ValueError: Unknown layer: LayerNormalization

How do I compile a BERT-keras model for TPU?

[Separius]

Hi @gaphex
I didn't write the TPU code and don't have a clue on what's going on.
HighCWu wrote it so I'll mention him.
@HighCWu can you help us?