using tfa

nmcardoso · Aug 22, 2023 · 85316a4 · 85316a4
1 parent 86ed7ad
commit 85316a4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/mergernet/estimators/base.py b/mergernet/estimators/base.py
@@ -3,6 +3,7 @@
 from typing import Callable, List, Tuple, Union
 
 import tensorflow as tf
+import tensorflow_addons as tfa
 
 from mergernet.core.constants import RANDOM_SEED
 from mergernet.core.experiment import Experiment
@@ -175,7 +176,8 @@ def get_dataaug_block(
 
   def get_metric(self, metric: str):
     if metric == 'f1':
-      return tf.keras.metrics.F1Score(name='f1', average='weighted')
+      return tfa.metrics.F1Score(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')
     elif metric == 'recall':

diff --git a/mergernet/jobs/__init__.py b/mergernet/jobs/__init__.py
@@ -1 +1,3 @@
+import tensorflow_addons as tfa
 
+tfa.register_all()