from tensorflow.contrib.eager.python import tfe

  data_dir = os.path.join(FLAGS.dir, "data")
  train_data = load_dataset(
      data_dir=data_dir, url=SOURCE_TRAIN_URL, batch_size=FLAGS.batch_size)
  eval_data = load_dataset(
      data_dir=data_dir, url=SOURCE_TEST_URL, batch_size=FLAGS.batch_size)

  model = RNNColorbot(
      rnn_cell_sizes=FLAGS.rnn_cell_sizes,
      label_dimension=3,
      keep_prob=FLAGS.keep_probability)
  optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.learning_rate)

  if FLAGS.no_gpu or tfe.num_gpus() <= 0:
    print(tfe.num_gpus())
    device = "/cpu:0"
  else:
    device = "/gpu:0"
  print("Using device %s." % device)

  log_dir = os.path.join(FLAGS.dir, "summaries")
  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(
