import tensorflow as tf

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