import tensorflow as tf

            run_config = tf.estimator.RunConfig(
                save_summary_steps=1,
                train_distribute=strategy,
                model_dir=FLAGS.output_dir,
                save_checkpoints_steps=FLAGS.save_ckpt_steps,
                log_step_count_steps=1,
            )
        else:
            distribution = tf.contrib.distribute.MirroredStrategy(
                num_gpus=FLAGS.num_gpus
            )
            run_config = tf.estimator.RunConfig(train_distribute=distribution)

    model_fn = model_fn_builder(
        bert_config=bert_config,
        init_checkpoint=FLAGS.init_checkpoint,
