| 162 | |
| 163 | |
| 164 | def split_state_dict_with_prefix(state_dict): |
| 165 | keys = sorted([key for key in state_dict if isinstance(key, str)]) |
| 166 | prefix_dict = {} |
| 167 | for key in keys: |
| 168 | prefix = key if "." not in key else key.split(".")[0] |
| 169 | if prefix not in prefix_dict: |
| 170 | prefix_dict[prefix] = [] |
| 171 | prefix_dict[prefix].append(key) |
| 172 | state_dicts = [] |
| 173 | for prefix, keys in prefix_dict.items(): |
| 174 | sub_state_dict = {key: state_dict[key] for key in keys} |
| 175 | state_dicts.append(sub_state_dict) |
| 176 | return state_dicts |
| 177 | |
| 178 | |
| 179 | def hash_state_dict_keys(state_dict, with_shape=True): |