Fix --end-learning-rate in polynomial LR schedule
Summary: Pull Request resolved:
@@ -60,7 +60,9 @@ 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
