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
    loss = None
    if (self.global_step + 1) % self.interval == 0:
        # Compute gradient at (w) per-GPU, and do not sync
        loss = closure(ddp_sync=False)  # type: ignore
        if loss:
            self.first_step()  # Compute e(w) and set weights to (w + (e(w)) separately per-GPU
            loss_dict = {}  # Dummy loss dict to ignore loss logging from w + e(w)
            if closure(loss_dict=loss_dict):  # type: ignore Compute gradient at (w + e(w))
                self.second_step()  # Reset weights to (w) and step base optimizer
            else:
                self.sub_e_w()  # If second forward-backward closure fails, reset weights to (w)
    else:
        loss = closure()
        if loss:
            self.base_optimizer.METHOD_NAME()
    self.global_step += 1
    return loss