Save optimizer state in Trainer #1416
Description
Is your feature request related to a problem? Please describe.
At the moment, we have the option to save model weights but not the optimizer state.
It could be useful to include this information if we'd like to resume training from a specific set of parameters (e.g. lr_schedule)!
Describe the solution you'd like
One solution could be to implement a .save method from the Trainer that also serializes the optimizer + training loop metadata.