native amp does not work when rewriting the optimizer_step function
🐛 Bug
If one rewrites optimizer_step like suggested in the docs (ie. https://pytorch-lightning.readthedocs.io/en/latest/optimizers.html#step-optimizers-at-arbitrary-intervals), the native AMP stops working. The issue happens due to the default optimizer_step performing a few native-amp-specific actions (see https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/core/lightning.py#L1176).
Please see this discussion for more detail.
https://pytorch-lightning.slack.com/archives/CRBLFHY79/p1588930218145400
Expected behavior
A naïve custom optimizer_step should work with native AMP.
I'm not sure how to really solve it. The simplest (but pretty bad in my opinion) option is to just write in the docs "If you want to use AMP and custom optimizer_step, add X, Y and Z to your code".
Additional context
The issue only happens on pytorch 1.6+ since that's where the native amp is implemented.