def METHOD_NAME(self, state_dict, strict=True):
    """Customized load."""
    self.language_model.METHOD_NAME(
        state_dict[self._language_model_key], strict=strict)
    if self.post_process and self.add_decoder:
        self.lm_head.METHOD_NAME(state_dict[self._lm_head_key],
                                     strict=strict)
    # Load word embeddings.
    if self.post_process and not self.pre_process and self.add_decoder:
        self.word_embeddings.METHOD_NAME(
            state_dict[self._word_embeddings_for_head_key], strict=strict)