(fullname)
| 7 | |
| 8 | |
| 9 | def convert_single_pth(fullname): |
| 10 | filename, ext = os.path.splitext(fullname) |
| 11 | checkpoint = torch.load(fullname, map_location='cpu', weights_only=True) |
| 12 | only_module = 'state_dict' not in checkpoint |
| 13 | state_dict = checkpoint if only_module else checkpoint['state_dict'] |
| 14 | torch.save(state_dict, fullname) |
| 15 | |
| 16 | if not only_module: |
| 17 | checkpoint.pop('state_dict') |
| 18 | fullname_trainer = filename + '_trainer_state' + ext |
| 19 | torch.save(checkpoint, fullname_trainer) |
| 20 | |
| 21 | |
| 22 | # This script is used to split pth files which generated before version 1.3.1 into two files. |
no test coverage detected
searching dependent graphs…