| 326 | self.optimizer = optimizer |
| 327 | |
| 328 | def state_dict(self): |
| 329 | # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT |
| 330 | model_state_dict, optimizer_state_dict = get_state_dict( |
| 331 | self.model, self.optimizer) |
| 332 | state_dict = { |
| 333 | 'model': model_state_dict, |
| 334 | 'optim': optimizer_state_dict, |
| 335 | } |
| 336 | if self.ema is not None: |
| 337 | ema_state_dict, _ = get_state_dict(self.ema, self.optimizer) |
| 338 | state_dict['ema'] = ema_state_dict |
| 339 | return state_dict |
| 340 | |
| 341 | def load_state_dict(self, state_dict): |
| 342 | # sets our state dicts on the model and optimizer, now that we've loaded |