diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index 5d601b549cdb886208a8f0f948b302b111b8b5d1..a65f86d163d81c968f8985d4b8604d1d1e36f398 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -723,11 +723,16 @@ class IcingIntensityNN:
         print('loss, acc: ', self.test_loss.result(), self.test_accuracy.result())
 
         labels = np.concatenate(self.test_labels)
+        self.test_labels = labels
+
         preds = np.concatenate(self.test_preds)
         self.test_probs = preds
-        preds = np.where(preds > 0.5, 1, 0)
 
-        self.test_labels = labels
+        if NumClasses == 2:
+            preds = np.where(preds > 0.5, 1, 0)
+        else:
+            preds = np.argmax(preds, axis=1)
+
         self.test_preds = preds
 
         self.h5f_tst.close()