| import numpy as np | |
| import torch | |
| import sys | |
| import os | |
| file = sys.argv[1] | |
| model = torch.load(file, map_location="cpu") | |
| if 'meta' in model.keys(): | |
| print("this file need not to convert.") | |
| exit(0) | |
| else: # this is a raw checkpoint | |
| meta_file = os.path.join(os.path.dirname(__file__), "Segmentation/example.pth") | |
| meta_data = torch.load(meta_file, map_location="cpu")['meta'] | |
| model = {'meta': meta_data, "state_dict": model} | |
| torch.save(model, file) | |
| print("converted to test-able file.") | |