diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index 552d533e043dce1a4d0f7612229f2d2006f894c5..e00c19e73107e5e4859abf4571b1f166c3fc395f 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -736,6 +736,7 @@ class IcingIntensityFCN:
         inputs = [mini_batch[0], mini_batch[1]]
         labels = mini_batch[2]
         pred = self.model(inputs, training=False)
+        pred = tf.reshape(pred, (pred.shape[0], NumLogits))
         t_loss = self.loss(labels, pred)
 
         self.test_labels.append(labels)