Transition smoothly after warmup in polynomial LR decay schedule

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

Differential Revision: D15318086

Pulled By: myleott

fbshipit-source-id: c6587737ca7b97edc97ad4aef5c5c9ac7e92b2f2

fairseq/optim/lr_scheduler/polynomial_decay_schedule.py

@@ -62,6 +62,7 @@ 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
             self.optimizer.set_lr(lr)
         return self.optimizer.get_lr()