import tensorflow as tf

        self.dataset = datasets.FlowersData(FLAGS.data_dir)
      else:
        raise ValueError('Unknown dataset. Must be one of imagenet or flowers.')

    self.local_parameter_device_flag = FLAGS.local_parameter_device
    if self.job_name:
      self.task_index = FLAGS.task_index
      self.cluster = tf.train.ClusterSpec({'ps': self.ps_hosts,
                                           'worker': self.worker_hosts})
      self.server = None

      if not self.server:
        self.server = tf.train.Server(self.cluster, job_name=self.job_name,
                                      task_index=self.task_index,
