Added AdamW from tf.keras.optimizers to external_configurables.

AdamW optimizer is supported in current version of tensorflow(https://www.tensorflow.org/api_docs/python/tf/keras/optimizers) PiperOrigin-RevId: 636486640
google · Jun 3, 2024 · 2bf410a · 2bf410a
1 parent 1507860
commit 2bf410a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/gin/tf/external_configurables.py b/gin/tf/external_configurables.py
@@ -98,6 +98,8 @@ def piecewise_constant(global_step, *args, **kwargs):
     tf.keras.optimizers.Adagrad, module='tf.keras.optimizers')
 config.external_configurable(
     tf.keras.optimizers.Adam, module='tf.keras.optimizers')
+config.external_configurable(
+    tf.keras.optimizers.AdamW, module='tf.keras.optimizers')
 config.external_configurable(
     tf.keras.optimizers.Adamax, module='tf.keras.optimizers')
 config.external_configurable(