(self, num_updates)
| 14 | self.step(0) |
| 15 | |
| 16 | def step(self, num_updates): |
| 17 | constant_lr = self.constant_lr |
| 18 | warmup = min(num_updates / self.warmup_updates, 1.0) |
| 19 | rsqrt_decay = max(self.warmup_updates, num_updates) ** -0.5 |
| 20 | rsqrt_hidden = self.hidden_size ** -0.5 |
| 21 | self.lr = max(constant_lr * warmup * rsqrt_decay * rsqrt_hidden, 1e-7) |
| 22 | for param_group in self.optimizer.param_groups: |
| 23 | param_group['lr'] = self.lr |
| 24 | return self.lr |
| 25 | |
| 26 | def get_lr(self): |
| 27 | return self.optimizer.param_groups[0]['lr'] |
no outgoing calls
no test coverage detected