| 63 | |
| 64 | # use _load_from_state_dict to enable checkpoint version control |
| 65 | def load(module, prefix=''): |
| 66 | # recursively check parallel module in case that the model has a |
| 67 | # complicated structure, e.g., nn.Module(nn.Module(DDP)) |
| 68 | if is_module_wrapper(module): |
| 69 | module = module.module |
| 70 | local_metadata = {} if metadata is None else metadata.get( |
| 71 | prefix[:-1], {}) |
| 72 | module._load_from_state_dict(state_dict, prefix, local_metadata, True, |
| 73 | all_missing_keys, unexpected_keys, |
| 74 | err_msg) |
| 75 | for name, child in module._modules.items(): |
| 76 | if child is not None: |
| 77 | load(child, prefix + name + '.') |
| 78 | |
| 79 | load(module) |
| 80 | load = None # break load->load reference cycle |