LR Scheduler instantiation fails when Optimizer loaded via DCP #315

The lr scheduler instantiation with an optimizer that has been loaded via distributed checkpointing (DCP) fails, as the inital_lr has not been set by the lr_scheduler (in the case of last_epoch != -1).

See here, for an extended discussion:
pytorch/pytorch#140900

Possible solution:
Torch Titan also checkpoints the lr scheduler:
https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/checkpoint.py#L266