(self, zero_grad=False)
| 31 | |
| 32 | @torch.no_grad() |
| 33 | def first_step(self, zero_grad=False): |
| 34 | grad_norm = self._grad_norm() |
| 35 | for group in self.param_groups: |
| 36 | scale = group["rho"] / (grad_norm + 1e-12) |
| 37 | |
| 38 | for p in group["params"]: |
| 39 | if p.grad is None: continue |
| 40 | e_w = p.grad * scale.to(p) |
| 41 | p.add_(e_w) # climb to the local maximum "w + e(w)" |
| 42 | self.state[p]["e_w"] = e_w |
| 43 | |
| 44 | if zero_grad: self.zero_grad() |
| 45 | |
| 46 | @torch.no_grad() |
| 47 | def second_step(self, zero_grad=False): |
no test coverage detected