MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / load_state_dict

Method load_state_dict

codegeex/oneflow/codegeex_model.py:895–908  ·  view source on GitHub ↗

Customized load.

(self, state_dict, strict=True)

Source from the content-addressed store, hash-verified

893 return state_dict_
894
895 def load_state_dict(self, state_dict, strict=True):
896 """Customized load."""
897
898 # Position embedding.
899 if self._top_query_embeddings_key in state_dict:
900 state_dict_ = state_dict[self._top_query_embeddings_key]
901 else:
902 # for backward compatibility.
903 state_dict_ = {}
904 for key in state_dict.keys():
905 if 'top_query_embeddings' in key:
906 state_dict_[key.split('top_query_embeddings.')[1]] \
907 = state_dict[key]
908 self.top_query_embeddings.load_state_dict(state_dict_, strict=strict)
909
910
911class TransformerLanguageModel(torch.nn.Module):

Callers

nothing calls this directly

Calls 1

load_state_dictMethod · 0.45

Tested by

no test coverage detected