| 14 | |
| 15 | |
| 16 | def process_checkpoint(in_file, out_file): |
| 17 | checkpoint = torch.load(in_file, map_location='cpu') |
| 18 | # remove optimizer for smaller file size |
| 19 | if 'optimizer' in checkpoint: |
| 20 | del checkpoint['optimizer'] |
| 21 | # if it is necessary to remove some sensitive data in checkpoint['meta'], |
| 22 | # add the code here. |
| 23 | torch.save(checkpoint, out_file) |
| 24 | sha = subprocess.check_output(['sha256sum', out_file]).decode() |
| 25 | final_file = out_file.rstrip('.pth') + '-{}.pth'.format(sha[:8]) |
| 26 | subprocess.Popen(['mv', out_file, final_file]) |
| 27 | |
| 28 | |
| 29 | def main(): |