(self, load_path)
| 321 | self.print_fn(f"model saved: {save_filename}") |
| 322 | |
| 323 | def load_model(self, load_path): |
| 324 | checkpoint = torch.load(load_path) |
| 325 | |
| 326 | self.model.load_state_dict(checkpoint['model']) |
| 327 | self.ema_model = deepcopy(self.model) |
| 328 | self.ema_model.load_state_dict(checkpoint['ema_model']) |
| 329 | self.optimizer.load_state_dict(checkpoint['optimizer']) |
| 330 | self.scheduler.load_state_dict(checkpoint['scheduler']) |
| 331 | self.it = checkpoint['it'] |
| 332 | self.print_fn('model loaded') |
| 333 | |
| 334 | |
| 335 | if __name__ == "__main__": |