model_fn.py

# 
# .... body of model_fn
#

optimizer = tf.train.AdamOptimizer()
  if FLAGS.use_tpu:
    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

  train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
  
#   return tf.estimator.EstimatorSpec(   # CPU or GPU estimator 
#     mode=mode,
#     loss=loss,
#     train_op=train_op,
#     predictions=predictions)

  return tf.contrib.tpu.TPUEstimatorSpec(  # TPU estimator 
      mode=mode,
      loss=loss,
      train_op=train_op,
      predictions=predictions)