From aec6d6edccf52c8ef029fb7cf711a98fb5117644 Mon Sep 17 00:00:00 2001
From: i4never <i4never@163.com>
Date: Tue, 22 Nov 2022 15:35:15 +0800
Subject: [PATCH] prevent sudden learning rate drop after warmup

---
 optimization.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/optimization.py b/optimization.py
index d33dabd91..41217e3f3 100644
--- a/optimization.py
+++ b/optimization.py
@@ -30,7 +30,7 @@ def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu):
 
   # Implements linear decay of the learning rate.
   learning_rate = tf.train.polynomial_decay(
-      learning_rate,
+      learning_rate * num_train_steps / (num_train_steps - num_warmup_steps),
       global_step,
       num_train_steps,
       end_learning_rate=0.0,