import tensorflow as tf


        if data_shape is None:
            self.data_shape = {i: s['shape'] for i, s in self.input_spec.items()}

        with tf.variable_scope('', reuse=tf.AUTO_REUSE):
            self._build_graph()

    def _gpu_tower(self, data, mode):
        # Split the batch between the GPUs (data parallelism)
        with tf.device('/cpu:0'):
            with tf.name_scope('{}_data_sharding'.format(mode)):
                batch_size = self.config['batch_size'] if (mode == Mode.TRAIN) \
                        else self.config['eval_batch_size']
                shards = {d: tf.unstack(v, num=batch_size*self.n_gpus, axis=0)
                          for d, v in data.items()}
                shards = [{d: tf.stack(v[i::self.n_gpus]) for d, v in shards.items()}
                          for i in range(self.n_gpus)]

