Fix --end-learning-rate in polynomial LR schedule

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

Differential Revision: D16068551

Pulled By: myleott

fbshipit-source-id: dddd8768b531032af7c4598af9dae3c6c00ff9ac

fairseq/optim/lr_scheduler/polynomial_decay_schedule.py

@@ -60,11 +60,13 @@ class PolynomialDecaySchedule(FairseqLRScheduler):
         """Update the learning rate after each update."""
         if self.args.warmup_updates > 0 and num_updates <= self.args.warmup_updates:
             self.warmup_factor = num_updates / float(self.args.warmup_updates)
-            self.optimizer.set_lr(self.warmup_factor * self.lr)
+            lr = self.warmup_factor * self.lr
+        elif num_updates >= self.total_num_update:
+            lr = self.end_learning_rate
         else:
             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)
+        self.optimizer.set_lr(lr)
         return self.optimizer.get_lr()