TF AdamWeightDecay fix for 2.11 (huggingface#20848)

* Fix incorrect import for the base optimizer for AdamWeightDecay * Fix incorrect import for the base optimizer for AdamWeightDecay
silverriver · Jan 6, 2023 · 9ede8c6 · 9ede8c6
1 parent 1678c53
commit 9ede8c6
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/transformers/optimization_tf.py b/src/transformers/optimization_tf.py
@@ -21,10 +21,10 @@
 import tensorflow as tf
 
 
-if hasattr(tf.keras, "optimizer") and hasattr(tf.keras.optimizer, "legacy"):
-    Adam = tf.keras.optimizer.legacy.Adam
-else:
-    Adam = tf.keras.optimizers.Adam
+try:
+    from tensorflow.keras.optimizers.legacy import Adam
+except ImportError:
+    from tensorflow.keras.optimizers import Adam
 
 
 class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):