TF AdamWeightDecay fix for 2.11 (#20848)

* Fix incorrect import for the base optimizer for AdamWeightDecay * Fix incorrect import for the base optimizer for AdamWeightDecay

TF AdamWeightDecay fix for 2.11 (#20848)
* Fix incorrect import for the base optimizer for AdamWeightDecay * Fix incorrect import for the base optimizer for AdamWeightDecay
d0bfdd20 · Matt · GitHub · d1d3ac94 · d0bfdd20
Unverified Commit d0bfdd20 authored Dec 20, 2022 by Matt Committed by GitHub Dec 20, 2022
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

src/transformers/optimization_tf.py src/transformers/optimization_tf.py +4 -4

No files found.
--- a/src/transformers/optimization_tf.py
+++ b/src/transformers/optimization_tf.py
@@ -21,10 +21,10 @@ from typing import Callable, List, Optional, Union
 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):