import tensorflow as tf

            # run_config = tf.estimator.RunConfig(
            #     experimental_distribute=tf.contrib.distribute.DistributeConfig(
            #         train_distribute=distribution,
            #         remote_cluster={
            #             'worker': ['localhost:5000', 'localhost:5001'],
            #         },
            #     )
            # )
            os.environ["TF_CONFIG"] = json.dumps(
                {
                    "cluster": {"worker": worker},
                    "task": {"type": "worker", "index": task_index},
                }
            )
            strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
            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)

