lr_scheduler.step() should pass epoch variable to be consistent with PyTorch _LRScheduler

[thinline72]

PyTorch _LRScheduler base class for schedulers supports passing epoch parameter to the lr_scheduler.step() function.

But when Lightning training loop updates learning rates, it doesn't pass the current epoch index. That leads to a wrong behaviour of those shedulers that relies on epoch param.

Some schedulers handle that under the hood, like CosineAnnealingWarmRestarts from PyTorch that increment last epoch index if epoch isn't provided. In my case I have a custom scheduler that updates learning rate after each batch via "interval": "step" and it also uses current epoch index. I can create a workaround for sure, but thought it'd be nice to get it fixed in the Pytorch Lightning too. Plus it wasn't obvious for me until I logged lr into tensorboard.

[github-actions]

Hi! thanks for your contribution!, great first issue!

[thinline72]

Actually, _LRScheduler also handles that by incrementing the last epoch index if epoch isn't provided.

Then it's an issue for all schedulers that update lr after each batch (like with "interval": "step") and rely on self.last_epoch. As it increments self.last_epoch after each batch, it basically becomes last step index instead of last epoch index.

[thinline72]

Oh, I see that PyTorch has a deprecation warning for the epoch param now. Although they still increment self.last_epoch at each step.

Anyway, probably it doesn't make sense to fix it, I'll close the issue. At least it'll be available in search if someone get stuck like me today :)