Loads the model weights from the specified path.
(
path: Path,
add_prefix: bool = False,
remove_prefix: bool = False
)
| 561 | |
| 562 | |
| 563 | def load_weights( |
| 564 | path: Path, |
| 565 | add_prefix: bool = False, |
| 566 | remove_prefix: bool = False |
| 567 | ) -> Dict[str, torch.Tensor]: |
| 568 | """ |
| 569 | Loads the model weights from the specified path. |
| 570 | """ |
| 571 | assert not (remove_prefix and add_prefix), \ |
| 572 | 'Cannot add and remove model prefix to the state dict simultaneously.' |
| 573 | |
| 574 | state = torch.load(path) |
| 575 | |
| 576 | # Handle cases where state_dict has 'model' key |
| 577 | state = state.get('model', state) |
| 578 | |
| 579 | if remove_prefix: |
| 580 | remove_model_prefix(state) |
| 581 | |
| 582 | if add_prefix: |
| 583 | add_model_prefix(state) |
| 584 | |
| 585 | return state |
no test coverage detected