| 29 | |
| 30 | # use _load_from_state_dict to enable checkpoint version control |
| 31 | def load(module, prefix=""): |
| 32 | # recursively check parallel module in case that the model has a |
| 33 | # complicated structure, e.g., nn.Module(nn.Module(DDP)) |
| 34 | local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) |
| 35 | module._load_from_state_dict( |
| 36 | state_dict, prefix, local_metadata, True, all_missing_keys, unexpected_keys, err_msg |
| 37 | ) |
| 38 | for name, child in module._modules.items(): |
| 39 | if child is not None: |
| 40 | load(child, prefix + name + ".") |
| 41 | |
| 42 | load(module) |
| 43 | load = None # break load->load reference cycle |