| 395 | raise RuntimeError('If NaN or Inf at iter 0, try lower lr. Otherwise please contact zhouyucong for a bug fix') |
| 396 | |
| 397 | def auto_denan_save(self): |
| 398 | if self.last_save_iter < 100 or self.tmp.current_step - self.last_save_iter > 100: |
| 399 | self.last_state_dict = {} |
| 400 | self.last_optim_state_dict = {} |
| 401 | # model state |
| 402 | for k,v in self.model.state_dict().items(): |
| 403 | self.last_state_dict[k] = v.cpu() |
| 404 | # optimizer state |
| 405 | self.last_optim_state_dict['state'] = {k:{'momentum_buffer':v['momentum_buffer'].cpu()} for k,v in self.optimizer.state_dict()['state'].items()} |
| 406 | #self.last_optim_state_dict['param_groups'] = copy.deepcopy(self.optimizer.state_dict()['param_groups']) # currently this is not needed |
| 407 | |
| 408 | self.last_save_iter = self.tmp.current_step |
| 409 | |
| 410 | def gather_result(self): |
| 411 | tmp = self.tmp |