include num_classes

nmcardoso · Aug 23, 2023 · 1937843 · 1937843
1 parent b455bcf
commit 1937843
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mergernet/estimators/base.py b/mergernet/estimators/base.py
@@ -176,7 +176,11 @@ def get_dataaug_block(
 
   def get_metric(self, metric: str):
     if metric == 'f1':
-      return tfa.metrics.F1Score(name='f1', average='weighted')
+      return tfa.metrics.F1Score(
+        num_classes=self.dataset.config.n_classes,
+        name='f1',
+        average='weighted'
+      )
       # return tf.keras.metrics.F1Score(name='f1', average='weighted') # tf 2.13
     elif metric == 'precision':
       return tf.keras.metrics.Precision(name='precision')