(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs)
| 3 | |
| 4 | class SAM(torch.optim.Optimizer): |
| 5 | def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs): |
| 6 | assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}" |
| 7 | |
| 8 | defaults = dict(rho=rho, adaptive=adaptive, **kwargs) |
| 9 | super(SAM, self).__init__(params, defaults) |
| 10 | |
| 11 | self.base_optimizer = base_optimizer(self.param_groups, **kwargs) |
| 12 | self.param_groups = self.base_optimizer.param_groups |
| 13 | self.defaults.update(self.base_optimizer.defaults) |
| 14 | |
| 15 | @torch.no_grad() |
| 16 | def first_step(self, zero_grad=False): |
nothing calls this directly
no outgoing calls
no test coverage detected