Pass second-order closure to all optimizers (not just LBFGS)
🚀 Feature
I could be wrong, but I noticed the following in the code of lightning module's optimizer_step
        if on_tpu:
            xm.optimizer_step(optimizer)
        elif using_native_amp:
            self.trainer.scaler.step(optimizer)
        elif using_lbfgs:
            optimizer.step(second_order_closure)
        else:
            optimizer.step()
If someone uses a custom optimizer that needs the closure returning the loss multiple times, it won't work.
Pitch
Since all classes that inherit from torch.optim.Optimizer have the step method accept a closure (even if they don't need it), we could just do
        if on_tpu:
            xm.optimizer_step(optimizer)
        elif using_native_amp:
            self.trainer.scaler.step(optimizer)
        # elif using_lbfgs:
        #     optimizer.step(second_order_closure)
        else:
            optimizer.step(second_order_closure)
and drop the "using_lbfgs" argument?
Alternatives
The user has to override the optimizer_step themself.