delay decay schedule until the end of warmup (#4940)

c9163a8d · Ali Modarressi · GitHub · f216b606 · c9163a8d
Unverified Commit c9163a8d authored Jun 24, 2020 by Ali Modarressi Committed by GitHub Jun 24, 2020
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/src/transformers/optimization_tf.py
+++ b/src/transformers/optimization_tf.py
@@ -44,7 +44,7 @@ class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):
            return tf.cond(
                global_step_float < warmup_steps_float,
                lambda: warmup_learning_rate,
-                lambda: self.decay_schedule_fn(step),
+                lambda: self.decay_schedule_fn(step - self.warmup_steps),
                name=name,
            )