How does LightningLite handle the grad scaler state dict of torch.amp?
Is the grad scaler included in the model.state_dict() after steup by model, optimizer = self.setup(model, optimizer)?
If not, how can I save and load the state of the grad scaler for resume?
Hey @hiyyg.
You can access the scale through the precision plugin as follows:
self._precision_plugin.scaler
I believe you could get the state and reload it manually.
Best,
T.C