| 388 | self.print_fn(f"model saved: {save_filename}") |
| 389 | |
| 390 | def load_model(self, load_path): |
| 391 | checkpoint = torch.load(load_path) |
| 392 | |
| 393 | self.model.load_state_dict(checkpoint['model']) |
| 394 | self.ema_model = deepcopy(self.model) |
| 395 | self.ema_model.load_state_dict(checkpoint['ema_model']) |
| 396 | self.optimizer.load_state_dict(checkpoint['optimizer']) |
| 397 | self.scheduler.load_state_dict(checkpoint['scheduler']) |
| 398 | self.it = checkpoint['it'] |
| 399 | self.print_fn('model loaded') |
| 400 | |
| 401 | def interleave_offsets(self, batch, nu): |
| 402 | groups = [batch // (nu + 1)] * (nu + 1) |