Commit 4647a7aa authored by Chen Qian's avatar Chen Qian Committed by A. Unique TensorFlower
Browse files

Internal change

PiperOrigin-RevId: 464672401
parent f8946f24
...@@ -133,9 +133,9 @@ class RankingTask(base_task.Task): ...@@ -133,9 +133,9 @@ class RankingTask(base_task.Task):
decay_steps=lr_config.decay_steps, decay_steps=lr_config.decay_steps,
decay_start_steps=lr_config.decay_start_steps) decay_start_steps=lr_config.decay_start_steps)
dense_optimizer = tf.keras.optimizers.Adam() dense_optimizer = tf.keras.optimizers.legacy.Adam()
embedding_optimizer = tf.keras.optimizers.get( embedding_optimizer = tf.keras.optimizers.get(
self.optimizer_config.embedding_optimizer) self.optimizer_config.embedding_optimizer, use_legacy_optimizer=True)
embedding_optimizer.learning_rate = lr_callable embedding_optimizer.learning_rate = lr_callable
feature_config = _get_tpu_embedding_feature_config( feature_config = _get_tpu_embedding_feature_config(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment