minor...

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
+
+        if NumClasses == 2:
             preds = np.where(preds > 0.5, 1, 0)
+        else:
+            preds = np.argmax(preds, axis=1)
 
-        self.test_labels = labels
         self.test_preds = preds
 
         self.h5f_tst.close()