(self, params, base_optimizer, rho=0.05, **kwargs)
| 22 | |
| 23 | class SAM(torch.optim.Optimizer): |
| 24 | def __init__(self, params, base_optimizer, rho=0.05, **kwargs): |
| 25 | assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}" |
| 26 | |
| 27 | defaults = dict(rho=rho, **kwargs) |
| 28 | super(SAM, self).__init__(params, defaults) |
| 29 | |
| 30 | self.base_optimizer = base_optimizer(self.param_groups, **kwargs) |
| 31 | self.param_groups = self.base_optimizer.param_groups |
| 32 | |
| 33 | @torch.no_grad() |
| 34 | def first_step(self, zero_grad=False): |
nothing calls this directly
no outgoing calls
no test coverage detected