Customized load.
(self, state_dict, strict=True)
| 821 | return state_dict_ |
| 822 | |
| 823 | def load_state_dict(self, state_dict, strict=True): |
| 824 | """Customized load.""" |
| 825 | |
| 826 | # Word embedding. |
| 827 | if self._word_embeddings_key in state_dict: |
| 828 | state_dict_ = state_dict[self._word_embeddings_key] |
| 829 | else: |
| 830 | # for backward compatibility. |
| 831 | state_dict_ = {} |
| 832 | for key in state_dict.keys(): |
| 833 | if 'word_embeddings' in key: |
| 834 | state_dict_[key.split('word_embeddings.')[1]] \ |
| 835 | = state_dict[key] |
| 836 | state_dict_["weight"] = state_dict_["weight"][:self.vocab_size] |
| 837 | self.word_embeddings.load_state_dict(state_dict_, strict=strict) |
| 838 | |
| 839 | # Position embedding. |
| 840 | if self._position_embeddings_key in state_dict: |
| 841 | state_dict_ = state_dict[self._position_embeddings_key] |
| 842 | else: |
| 843 | # for backward compatibility. |
| 844 | state_dict_ = {} |
| 845 | for key in state_dict.keys(): |
| 846 | if 'position_embeddings' in key: |
| 847 | state_dict_[key.split('position_embeddings.')[1]] \ |
| 848 | = state_dict[key] |
| 849 | self.position_embeddings.load_state_dict(state_dict_, strict=strict) |
| 850 | |
| 851 | |
| 852 | class QueryEmbedding(torch.nn.Module): |
no outgoing calls
no test coverage detected