diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index 36a6c57f141ef914098816fb8068f7e8e4f4c6f2..97ee728d01a0a875a617fab2848ed67366e2e7ba 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -381,8 +381,10 @@ class IcingIntensityNN:
         fc = tf.keras.layers.BatchNormalization()(fc)
         print(fc.shape)
 
-        # activation = tf.nn.softmax # For multi-class
-        activation = tf.nn.sigmoid  # For binary
+        if NumClasses == 2:
+            activation = tf.nn.sigmoid  # For binary
+        else:
+            activation = tf.nn.softmax  # For multi-class
 
         # Called logits, but these are actually probabilities see activation
         logits = tf.keras.layers.Dense(NumLogits, activation=activation)(fc)