(model: jt.Module, ckpt_dir, file_weight_map)
| 46 | jt.gc() |
| 47 | |
| 48 | def check_state_dict(model: jt.Module, ckpt_dir, file_weight_map): |
| 49 | for filename, names in file_weight_map.items(): |
| 50 | cur_state_dict = torch.load(os.path.join(ckpt_dir, filename)) |
| 51 | for name in names: |
| 52 | assert np.equal( |
| 53 | model.state_dict()[name].numpy(), cur_state_dict[name].numpy()).all() |
| 54 | |
| 55 | # gc to reduce memory usage |
| 56 | del cur_state_dict |
| 57 | jt.sync_all() |
| 58 | jt.gc() |
nothing calls this directly
no outgoing calls
no test coverage detected