(self, load_dir)
| 23 | torch.save(self.model.state_dict(), save_path) |
| 24 | |
| 25 | def from_pretrained(self, load_dir): |
| 26 | state_dict = torch.load(load_dir, map_location='cpu') |
| 27 | if 'model' in state_dict: |
| 28 | state_dict=state_dict['model'] |
| 29 | state_dict={k[6:]:v for k,v in state_dict.items() if k.startswith('model.')} |
| 30 | # for k in self.model.state_dict(): |
| 31 | # if k not in state_dict: |
| 32 | # assert k[:-2] in state_dict |
| 33 | # state_dict[k]=state_dict.pop(k[:-2]) |
| 34 | state_dict = align_and_update_state_dicts(self.model.state_dict(), state_dict) |
| 35 | self.model.load_state_dict(state_dict, strict=False) |
| 36 | return self |
no test coverage detected