(param, step, exp_avg, exp_avg_sq, beta1, beta2, lr, eps)
| 88 | return exp_avg_sq.mul(beta2).addcmul(grad, grad, value=1 - beta2) |
| 89 | |
| 90 | def update_param(param, step, exp_avg, exp_avg_sq, beta1, beta2, lr, eps): |
| 91 | bias_correction1 = 1 - torch.pow(beta1, step) |
| 92 | bias_correction2 = (1 - torch.pow(beta2, step)).sqrt() |
| 93 | step_size = (lr / bias_correction1).neg() |
| 94 | denom = (exp_avg_sq.sqrt() / (bias_correction2 * step_size)).add(eps / step_size) |
| 95 | return torch.add(param, torch.div(exp_avg, denom)) |
| 96 | |
| 97 | # Our full Adam implementation |
| 98 | def foreach_map_adam( |
nothing calls this directly
no outgoing calls
no test coverage detected