import tensorflow as tf

    def _pred_graph(self, data):
        with tf.name_scope('pred'):
            with tf.device('/gpu:0'):
                pred_out = self._model(data, Mode.PRED, **self.config)
        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
