REALM Training on TPU

Hi, I'm trying to train REALM on Colab Free TPU with my own dataset, pre-trained model and modified code.

My question is : Can I train REALM on TPU (with GRPC call for dynamic dataset).

I deployed GRPC server on Compute Engine VM, and confirmed I can make a GRPC call with following code.
"aaa.aaa.aaa.aaa" is VM's external IP, "bbbb" is the port of GRPC server.
```
import sys
sys.path.append("./language")
from tensorflow.compat import v1 as tf
from language.realm.preprocessing import _make_rpc_op

tf.enable_eager_execution()

with tf.Session() as sess:
result = sess.run([_make_rpc_op("aaa.aaa.aaa.aaa:bbbb", 1000)])
print(result[0].response)

# b'\n\xbb\x94\x01\n\x8c\x18\n\x016....
```
I was able to parse the received serialized example to the features. so my GRPC server seems working well.

But when I invoked my training loop with `estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps)`, training seems get stuck after following log messages.(I removed some duplicate messages manually.)
```
...
INFO:tensorflow:Create CheckpointSaverHook.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:TPU job name worker
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Saving checkpoints for 0 into gs://somewhere/realm/model/model.ckpt.
WARNING:tensorflow:From /tensorflow-1.15.2/python3.6/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py:751: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer Variable.assign which has equivalent behavior in 2.X.
INFO:tensorflow:Initialized dataset iterators in 1 seconds
INFO:tensorflow:Installing graceful shutdown hook.
INFO:tensorflow:Creating heartbeat manager for ['/job:worker/replica:0/task:0/device:CPU:0']
INFO:tensorflow:Configuring worker heartbeat: shutdown_mode: WAIT_FOR_COORDINATOR
INFO:tensorflow:Init TPU system
INFO:tensorflow:Initialized TPU in 5 seconds
INFO:tensorflow:Starting infeed thread controller.
INFO:tensorflow:Starting outfeed thread controller.
INFO:tensorflow:Enqueue next (10) batch(es) of data to infeed.
INFO:tensorflow:Dequeue next (10) batch(es) of data from outfeed.
```
Training seems running, but no log output and checkpoints show up after that.
Currently, I used Colab Free TPU for now, and `train_preprocessing_servers` was specified as follows.
```
FLAGS.train_preprocessing_servers
# ["aaa.aaa.aaa.aaa:bbbb"]
```
I also tried with VM's internal ip address, but got no luck.

After that, I wrote serialized examples to TFRecord format file. and replaced `get_dynamic_dataset()` with following code.
```
def get_dynamic_dataset(preprocessing_servers,
featurizer,
num_input_threads):

def _parse_example(serialized_example):
return nest_utils.tf_example_to_structure(
serialized_example, featurizer.query_and_docs_feature_structure)

realm_examples = "gs://somewhere/realm_examples.tfr"
dataset = tf.data.TFRecordDataset([realm_examples], compression_type="GZIP")
return dataset.map(_parse_example, num_parallel_calls=num_input_threads)

from language.realm import model
model.get_dynamic_dataset = get_dynamic_dataset
```

With only this modification, training works well, so there seems to be a problem with GRPC call.
I want to train with dynamic dataset, is there a way to achieve that with TPU?