(module, local_state_dict, prefix='')
| 610 | state_dict._metadata = metadata # type: ignore[attr-defined] |
| 611 | |
| 612 | def load(module, local_state_dict, prefix=''): |
| 613 | local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) |
| 614 | if assign: |
| 615 | local_metadata['assign_to_params_buffers'] = assign |
| 616 | _zero3_load_from_state_dict(module, |
| 617 | local_state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs) |
| 618 | for name, child in module._modules.items(): |
| 619 | if child is not None: |
| 620 | child_prefix = prefix + name + '.' |
| 621 | child_state_dict = {k: v for k, v in local_state_dict.items() if k.startswith(child_prefix)} |
| 622 | load(child, child_state_dict, child_prefix) # noqa: F821 |
| 623 | |
| 624 | # Note that the hook can modify missing_keys and unexpected_keys. |
| 625 | incompatible_keys = _IncompatibleKeys(missing_keys, unexpected_keys) |
| 626 | for hook in module._load_state_dict_post_hooks.values(): |
| 627 | out = hook(module, incompatible_keys) |
| 628 | assert out is None, ( |
| 629 | "Hooks registered with ``register_load_state_dict_post_hook`` are not" |
| 630 | "expected to return new values, if incompatible_keys need to be modified," |
| 631 | "it should be done inplace." |
| 632 | ) |
| 633 | |
| 634 | load(self, state_dict) |
| 635 | del load |
no test coverage detected