Torch optimisers expect the step function to run the closure immediately if passed

[tomogwen]

As per title. If the closure is not None, we should assume that the parameter's gradients have not been computed, and immediately run the closure. See, e.g., the step function in torch.optim.adamw as an example of this.

Without this, the library won't work out the box with PyTorch Lightning, and I imagine this is also the source of the problem for #90 with the HF transformers/accelerate libraries.

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.