Fix bug with saving optimizer states with MonolithicCheckpointSaver Callback (#310)
* Fix bug with saving optimizer states with mono ckpt saver
* lint
---------
Co-authored-by: Abhi Venigalla <77638579+abhi-mosaic@users.noreply.github.com>
Co-authored-by: Vitaliy Chiley <6439018+vchiley@users.noreply.github.com>
Co-authored-by: root <vitaliy@mosaicml.com>