Use LRScheduler instead of _LRScheduler

[janeyx99]

Update code to be compatible with LRScheduler too

I recently exposed LRScheduler to be public (instead of _LRScheduler) in pytorch/pytorch#88503. I see several instance checks in this current repo that use _LRScheduler, and I believe these should be updated to be LRScheduler.

Why?

This would be a good step in the right direction. Also, downstream users who subscribe to new LRScheduler (and not _LRScheduler) now will run into issues.

Alternatives

Maybe the best thing to do for forward and backward compatibility is to accept LRScheduler whenever there is _LRScheduler too.

[vfdev-5]

@janeyx99 thanks for the update! Indeed this change broke our CI running with pytorch-nightly :)