import tensorflow as tf

  config = tf.ConfigProto()
  if FLAGS.xla:
    config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
  if FLAGS.use_hvd:
    config.gpu_options.visible_device_list = str(hvd.local_rank())
    config.gpu_options.allow_growth=True

  run_config = tf.estimator.RunConfig(
      model_dir=FLAGS.output_dir,
      keep_checkpoint_max=FLAGS.keep_checkpoint_max,
      save_checkpoints_steps=FLAGS.save_checkpoints_steps,
      log_step_count_steps=FLAGS.hooking_frequence,
      session_config=config)

