from torch import load, save

if __name__ == '__main__':
    print("优化器通常不会被用于推理阶段，如果只用于推理可以去除优化器以减小模型体积\n")
    input_path = input("请输入模型的路径：")
    output_path = f"{input_path.split('.')[0]}_inference.pth"
    checkpoint_dict = load(input_path, map_location='cpu')
    checkpoint_dict_new = {}
    for k, v in checkpoint_dict.items():
        if k == "optimizer":
            print(f"remove optimizer")
            continue
        checkpoint_dict_new[k] = v
    save(checkpoint_dict_new, output_path)
    print("finish")
    print(output_path)