hub version elmo with colab tpu runtime produces eorror

I'm using the hub version elmo model to do the embedding, like below:

`sess=tf.Session()

K.set_session(sess)

elmo_model = c("https://tfhub.dev/google/elmo/2", trainable=True)

trainX=elmo_model(inputs={"tokens": X_tr,"sequence_len": len_X_tr
},
signature="tokens",
as_dict=True)["elmo"]

sess.run(tf.global_variables_initializer())

sess.run(tf.tables_initializer())

embed_train=sess.run(trainX)`

when i run with GPU runtime type, there's no problem. but when i change the runtime to tpu, it produces errors as following:
InvalidArgumentError (see above for traceback): seq_lens(189) > input.dims(1)
[[node module_apply_tokens/bilm/ReverseSequence (defined at /usr/local/lib/python3.6/dist-packages/tensorflow_hub/native_module.py:547) ]]

Could anyone help me with this? Thanks a lot!