import tensorflow as tf

    'especially for TPU inference.')
flags.DEFINE_string('model_name', 'amoeba_net',
                    'Serving model name used for the model server.')
flags.DEFINE_multi_integer(
    'inference_batch_sizes', [8],
    'Known inference batch sizes used to warm up for each core.')

FLAGS = flags.FLAGS


def build_run_config():
  """Return RunConfig for TPU estimator."""
  tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
      FLAGS.tpu,
      zone=FLAGS.tpu_zone,
      project=FLAGS.gcp_project)

  eval_steps = FLAGS.num_eval_images // FLAGS.eval_batch_size
  iterations_per_loop = (eval_steps if FLAGS.mode == 'eval'
                         else FLAGS.iterations_per_loop)
  save_checkpoints_steps = FLAGS.save_checkpoints_steps or iterations_per_loop
  run_config = tf.contrib.tpu.RunConfig(
      cluster=tpu_cluster_resolver,
      model_dir=FLAGS.model_dir,
