| 465 | self.optimizer.step() |
| 466 | |
| 467 | def _save(self, step): |
| 468 | real_model = self.model |
| 469 | # real_generator = (self.generator.module |
| 470 | # if isinstance(self.generator, torch.nn.DataParallel) |
| 471 | # else self.generator) |
| 472 | |
| 473 | model_state_dict = real_model.state_dict() |
| 474 | # generator_state_dict = real_generator.state_dict() |
| 475 | checkpoint = { |
| 476 | 'models': model_state_dict, |
| 477 | # 'generator': generator_state_dict, |
| 478 | 'opt': self.args, |
| 479 | 'optim': self.optimizer, |
| 480 | } |
| 481 | checkpoint_path = os.path.join(self.args.model_path, 'model_step_%d.pt' % step) |
| 482 | logger.info("Saving checkpoint %s" % checkpoint_path) |
| 483 | # checkpoint_path = '%s_step_%d.pt' % (FLAGS.model_path, step) |
| 484 | if not os.path.exists(checkpoint_path): |
| 485 | torch.save(checkpoint, checkpoint_path) |
| 486 | return checkpoint, checkpoint_path |
| 487 | |
| 488 | def _start_report_manager(self, start_time=None): |
| 489 | """ |