diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index 9002f2902b2a5151e31dd00c32e6a8300b4852f3..5f0ba197764e6c3f6806b8981f7ecaeeedb5985d 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -627,7 +627,7 @@ class IcingIntensityFCN:
             activation = tf.nn.softmax  # For multi-class
 
         # Called logits, but these are actually probabilities, see activation
-        logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, activation=activation)(conv)
+        logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, name='probability', activation=activation)(conv)
 
         print(logits.shape)