(ckpt_path, location='cpu', exclude_buffers=None)
| 21 | # return state_dict |
| 22 | |
| 23 | def load_state_dict(ckpt_path, location='cpu', exclude_buffers=None): |
| 24 | _, extension = os.path.splitext(ckpt_path) |
| 25 | if extension.lower() == ".safetensors": |
| 26 | import safetensors.torch |
| 27 | state_dict = safetensors.torch.load_file(ckpt_path, device=location).to(torch.float16) |
| 28 | else: |
| 29 | state_dict = get_state_dict(torch.load(ckpt_path, map_location=torch.device(location))) |
| 30 | |
| 31 | if exclude_buffers: |
| 32 | state_dict = {k: v for k, v in state_dict.items() if not any(buf_name in k for buf_name in exclude_buffers)} |
| 33 | |
| 34 | print(f'Loaded state_dict from [{ckpt_path}]') |
| 35 | # for name, param in state_dict.items(): |
| 36 | # print(f"Layer: {name}, Shape: {param.shape}") |
| 37 | return state_dict |
| 38 | |
| 39 | |
| 40 |
no test coverage detected