import tensorflow as tf

  train_steps_per_epoch = int(
      math.ceil(hparams.num_train_images / float(hparams.train_batch_size)))
  eval_steps = hparams.num_eval_images // hparams.eval_batch_size
  eval_batch_size = (None if mode == 'train' else
                     hparams.eval_batch_size)

  model = model_lib.AmoebaNetEstimatorModel(hparams, model_dir)

  if hparams.use_tpu:
    run_config = build_run_config()
    image_classifier = tf.contrib.tpu.TPUEstimator(
        model_fn=model.model_fn,
        use_tpu=True,
        config=run_config,
        params=estimator_parmas,
        predict_batch_size=eval_batch_size,
        train_batch_size=hparams.train_batch_size,
        eval_batch_size=eval_batch_size,
        export_to_tpu=FLAGS.export_to_tpu,
        experimental_exported_model_uses_all_cores=FLAGS
