How does LightningLite handle the grad scaler state dict of torch.amp?

[hiyyg]

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?

[tchaton]

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

[hiyyg]

OK. I will try that.