import tensorflow as tf

        max_seq_length=FLAGS.max_seq_length,
        max_predictions_per_seq=FLAGS.max_predictions_per_seq,
        is_training=True,
        batch_size=FLAGS.train_batch_size,
        use_hvd=FLAGS.use_hvd)

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

    estimator.train(input_fn=train_input_fn, max_steps=FLAGS.num_train_steps, hooks=hooks)

  if FLAGS.do_eval:
    tf.logging.info("***** Running evaluation *****")
    tf.logging.info("  Batch size = %d", FLAGS.eval_batch_size)
