diff --git a/sciencebeam_gym/trainer/models/pix2pix/pix2pix_model.py b/sciencebeam_gym/trainer/models/pix2pix/pix2pix_model.py
index 3e0aed22f0eeaf9c7d6a7aed3d46f1398e35e572..c3b0d52ab36d1eec36c5967485b2d6acc1958411 100644
--- a/sciencebeam_gym/trainer/models/pix2pix/pix2pix_model.py
+++ b/sciencebeam_gym/trainer/models/pix2pix/pix2pix_model.py
@@ -382,6 +382,7 @@ class Model(object):
     else:
       with tf.name_scope('evaluation'):
         if self.dimension_colors:
+          tensors.output_layer_labels = tf.constant(self.dimension_labels)
           colors_tensor = tf.constant(
             self.dimension_colors_with_unknown,
             dtype=tf.float32