diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index 5409e32c613073c355a03d1f10e8249d66df6c68..39b7f5883cacd9b8108d9b51fa885baee679fb03 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -463,14 +463,14 @@ class IcingIntensityNN:
             self.train_accuracy = tf.keras.metrics.BinaryAccuracy(name='train_accuracy')
             self.test_accuracy = tf.keras.metrics.BinaryAccuracy(name='test_accuracy')
             self.test_auc = tf.keras.metrics.AUC(name='test_auc')
-            self.test_f1 = tfa.metrics.F1Score(name='test_f1')
+            self.test_f1 = tfa.metrics.F1Score(NumClasses, name='test_f1')
             self.test_recall = tf.keras.metrics.Recall(name='test_recall')
             self.test_precision = tf.keras.metrics.Precision(name='test_precision')
         else:
             self.train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
             self.test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')
             self.test_auc = tf.keras.metrics.AUC(name='test_auc')
-            self.test_f1 = tfa.metrics.F1Score(name='f1_score')
+            self.test_f1 = tfa.metrics.F1Score(NumClasses, name='f1_score')
             self.test_recall = tf.keras.metrics.Recall(name='test_recall')
             self.test_precision = tf.keras.metrics.Precision(name='test_precision')