simplify code

    def word_embeddings_weight(self):

        if hasattr(self.language_model, 'embedding') and \

                self.language_model.embedding is not None:

        if self.pre_process:

            return self.language_model.embedding.word_embeddings.weight

        else:

            if not self.share_word_embeddings:

        #    the two word_embeddings layers to ensure that every applied weight

        #    update is the same on both stages.

        if mpu.is_pipeline_last_stage() and \

                (not hasattr(self.language_model, 'embedding') or 

                 self.language_model.embedding is None):

                not self.pre_process:

            assert not mpu.is_pipeline_first_stage()

            self._word_embeddings_for_head_key = 'word_embeddings_for_head'

            # set word_embeddings weights to 0 here, then copy first

        # Zero out initial weights for decoder embedding.

        # NOTE: We don't currently support T5 with the interleaved schedule.

        if not mpu.is_pipeline_first_stage(ignore_virtual=True) and \

                mpu.is_rank_in_embedding_group() and \

                hasattr(self.language_model, 'embedding') and \

                self.language_model.embedding is not None:

                self.pre_process:

            self.language_model.embedding.zero_parameters()

        # Ensure that first and last stages have the same initial parameter