(self, epoch)
| 113 | self.scheduler.step() |
| 114 | |
| 115 | def save_checkpoint(self, epoch): |
| 116 | state_dict = {"state_dict": self.model.state_dict(), |
| 117 | "optimizer": self.optimizer.state_dict(), |
| 118 | "epoch": epoch} |
| 119 | torch.save(state_dict, |
| 120 | '{}/pth/model_{}.pth'.format(self.config.OUTPUT.DIR, epoch)) |
| 121 | |
| 122 | def load_checkpoint(self, weight_path, resume=False): |
| 123 | checkpoint = torch.load(weight_path) |