import tensorflow as tf

            )
        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,
