(self, closure=None)
| 56 | |
| 57 | @torch.no_grad() |
| 58 | def step(self, closure=None): |
| 59 | assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided" |
| 60 | closure = torch.enable_grad()(closure) # the closure should do a full forward-backward pass |
| 61 | |
| 62 | self.first_step(zero_grad=True) |
| 63 | closure() |
| 64 | self.second_step() |
| 65 | |
| 66 | def _grad_norm(self): |
| 67 | shared_device = self.param_groups[0]["params"][0].device # put everything on the same device, in case of model parallelism |
no test coverage detected