import tensorflow as tf

        batch_size=FLAGS.eval_batch_size,
        use_hvd=FLAGS.use_hvd)

    if FLAGS.auto_recover:
      hooks.append(tf.data.experimental.CheckpointInputPipelineHook(estimator))

    train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps, hooks=hooks)
    eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn, steps=FLAGS.max_eval_steps)
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)


if __name__ == "__main__":
  # flags.mark_flag_as_required("input_file")
  flags.mark_flag_as_required("bert_config_file")
  flags.mark_flag_as_required("output_dir")
  tf.app.run()
