(self, optimizer)
| 3 | |
| 4 | class RSQRTSchedule(object): |
| 5 | def __init__(self, optimizer): |
| 6 | super().__init__() |
| 7 | self.optimizer = optimizer |
| 8 | self.constant_lr = hparams['lr'] |
| 9 | self.warmup_updates = hparams['warmup_updates'] |
| 10 | self.hidden_size = hparams['hidden_size'] |
| 11 | self.lr = hparams['lr'] |
| 12 | for param_group in optimizer.param_groups: |
| 13 | param_group['lr'] = self.lr |
| 14 | self.step(0) |
| 15 | |
| 16 | def step(self, num_updates): |
| 17 | constant_lr = self.constant_lr |