def METHOD_NAME(self, closure=None):
    assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
    closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass
    enable_running_stats(self.model)
    loss = closure()
    self.first_step(zero_grad=True)
    disable_running_stats(self.model)
    closure()
    self.second_step()
    return loss