import tensorflow as tf

  tf.gfile.MakeDirs(log_dir)
  train_summary_writer = tf.contrib.summary.create_file_writer(
      os.path.join(log_dir, "train"), flush_millis=10000)
  test_summary_writer = tf.contrib.summary.create_file_writer(
      os.path.join(log_dir, "eval"), flush_millis=10000, name="eval")

  with tf.device(device):
    for epoch in range(FLAGS.num_epochs):
      start = time.time()
      with train_summary_writer.as_default():
        train_one_epoch(model, optimizer, train_data, FLAGS.log_interval)
      end = time.time()
