(self, closure=None)
| 39 | |
| 40 | @torch.no_grad() |
| 41 | def step(self, closure=None): |
| 42 | assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided" |
| 43 | closure = torch.enable_grad()(closure) # the closure should do a full forward-backward pass |
| 44 | |
| 45 | self.first_step(zero_grad=True) |
| 46 | closure() |
| 47 | self.second_step() |
| 48 | |
| 49 | def _grad_norm(self): |
| 50 | shared_device = self.param_groups[0]["params"][0].device # put everything on the same device, in case of model parallelism |
no test coverage detected