Fixed wrong lr initialization when loading checkpoints

[AleHD]

For some reason the lr in the parameter groups of the optimizers is set to zero in the very first iteration. This is problematic when loading checkpoints as it essentially means you are no longer following the same training as an uninterrupted run, but rather skipping one optimization iteration. This PR fixes this issue.

Minimalistic example:

Green follows a 10 iteration training without interruption. Orange is the current behaviour, loading from iteration 5 and quickly diverging because of the incorrect lr set when loading. Brown (overlapping green) is the fix loading from iteration 5 but now follows the same expected route.

[AleHD]

I have introduced Kyle's changes to finally fix the issue. Now it looks to be resuming succesfully regerdless of stage of the learning rate (during warmup, steady, cooldown or after cooldown all work fine). Should be good to merge :)