Fix warmup for polynomial decay schedule

Summary: Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/611

Differential Revision: D15541303

Pulled By: myleott

fbshipit-source-id: 279ca813437c834fca49576a48b75cbf1fdf0e76

fairseq/optim/lr_scheduler/polynomial_decay_schedule.py

@@ -62,7 +62,9 @@ class PolynomialDecaySchedule(FairseqLRScheduler):
             self.warmup_factor = num_updates / float(self.args.warmup_updates)
             self.optimizer.set_lr(self.warmup_factor * self.lr)
         else:
-            num_updates -= self.args.warmup_updates
-            lr = (self.lr - self.end_learning_rate) * (1 - num_updates / self.total_num_update) ** (self.power) + self.end_learning_rate
+            warmup = self.args.warmup_updates
+            lr_range = self.lr - self.end_learning_rate
+            pct_remaining = 1 - (num_updates - warmup) / (self.total_num_update - warmup)
+            lr = lr_range * pct_remaining ** (self.power) + self.end_learning_rate
             self.optimizer.set_lr(lr)
         return self.optimizer.get_lr()