import tensorflow as tf

        self.pred_out = {n: tf.identity(p, name=n) for n, p in pred_out.items()}

    def _build_graph(self):
        # Training and evaluation network, if tf datasets provided
        if self.datasets:
            # Generate iterators for the given tf datasets
            self.dataset_iterators = {}
            with tf.device('/cpu:0'):
                for n, d in self.datasets.items():
                    if n == 'training':
                        train_batch = self.config['batch_size']*self.n_gpus
                        d = d.repeat().batch(train_batch).prefetch(train_batch)
                        self.dataset_iterators[n] = d.make_one_shot_iterator()
                    else:
