(directory_path: str, device)
| 5 | import pickle |
| 6 | |
| 7 | def load_state(directory_path: str, device): |
| 8 | with open(directory_path + '/tree.pkl', 'rb') as f: |
| 9 | tree = pickle.load(f) |
| 10 | state = torch.load(directory_path + '/model_state.pth', map_location=device) |
| 11 | tree.load_state_dict(state) |
| 12 | return tree |
| 13 | |
| 14 | def init_tree(tree: ProtoTree, optimizer, scheduler, device, args: argparse.Namespace): |
| 15 | epoch = 1 |