(state_dicts, head_prefix,
expected_keys_without_head_prefix,
missing_keys)
| 174 | """ |
| 175 | |
| 176 | def _add_head_prefix_to_state_dict(state_dicts, head_prefix, |
| 177 | expected_keys_without_head_prefix, |
| 178 | missing_keys): |
| 179 | new_state_dict = OrderedDict() |
| 180 | for name, module in state_dicts.items(): |
| 181 | if name in expected_keys_without_head_prefix: |
| 182 | name_with_head = '.'.join([head_prefix, name]) |
| 183 | new_state_dict[name_with_head] = module |
| 184 | expected_keys_without_head_prefix.remove(name) |
| 185 | missing_keys = list(set(missing_keys) - set([name_with_head])) |
| 186 | else: |
| 187 | new_state_dict[name] = module |
| 188 | |
| 189 | missing_head_keys = [] |
| 190 | if len(expected_keys_without_head_prefix) > 0: |
| 191 | missing_head_keys = expected_keys_without_head_prefix.copy() |
| 192 | return new_state_dict, missing_head_keys, missing_keys |
| 193 | |
| 194 | def _find_mismatched_keys( |
| 195 | state_dicts, |
no test coverage detected
searching dependent graphs…