import tensorflow as tf

            use_tpu=FLAGS.use_tpu,
            model_fn=model_fn,
            config=run_config,
            train_batch_size=FLAGS.train_batch_size,
            eval_batch_size=FLAGS.eval_batch_size,
        )
    else:
        estimator = tf.estimator.Estimator(
            model_fn=model_fn,
            config=run_config,
            params={
                "batch_size": FLAGS.train_batch_size
                if FLAGS.do_train
                else FLAGS.eval_batch_size,
