def forward(
        self,
        tgt,
        memory_bank,
        state,
        memory_lengths=None,
        step=None,
        cache=None,
        memory_masks=None,
    ):
        """
        See :obj:`onmt.modules.RNNDecoderBase.forward()`
        """

        src_words = state.src
        tgt_words = tgt
        src_batch, src_len = src_words.size()
        tgt_batch, tgt_len = tgt_words.size()

        # Run the forward pass of the TransformerDecoder.
        # emb = self.embeddings(tgt, step=step)
        emb = self.embeddings(tgt)
        assert emb.dim() == 3  # len x batch x embedding_dim

        output = self.pos_emb(emb, step)

        src_memory_bank = memory_bank
        padding_idx = self.embeddings.padding_idx
        tgt_pad_mask = (
            tgt_words.data.eq(padding_idx)
            .type(torch.uint8)
            .unsqueeze(1)
            .expand(tgt_batch, tgt_len, tgt_len)
        )

        if not memory_masks is None:
            src_len = memory_masks.size(-1)
            src_pad_mask = memory_masks.expand(src_batch, tgt_len, src_len)

        else:
            src_pad_mask = (
                src_words.data.eq(padding_idx)
                .unsqueeze(1)
                .expand(src_batch, tgt_len, src_len)
            )

        if state.cache is None:
            saved_inputs = []

        for i in range(self.num_layers):
            prev_layer_input = None
            if state.cache is None:
                if state.previous_input is not None:
                    prev_layer_input = state.previous_layer_inputs[i]
            output, all_input = self.transformer_layers[i](
                output,
                src_memory_bank,
                src_pad_mask,
                tgt_pad_mask,
                previous_input=prev_layer_input,
                layer_cache=state.cache["layer_{}".format(i)]
                if state.cache is not None
                else None,
                step=step,
            )
            if state.cache is None:
                saved_inputs.append(all_input)

        if state.cache is None:
            saved_inputs = torch.stack(saved_inputs)

        output = self.layer_norm(output)

        # Process the result and update the attentions.

        if state.cache is None:
            state = state.update_state(tgt, saved_inputs)

        return output, state

        See :obj:`onmt.modules.RNNDecoderBase.forward()`
        
utils_nlp/models/transformers/bertsum/decoder.py