code: def _get_memory_states(
      self, decoder_states, decoder_mems_list=None, i=0
  ):
      if decoder_mems_list is not None:
          inp1 = torch.transpose(
              decoder_mems_list[i], 1, 2
          )  # Putting seq_len to last dim to handle export cases
          inp2 = torch.transpose(decoder_states, 1, 2)
          memory_states = torch.cat((inp1, inp2), dim=2)
          memory_states = torch.transpose(
              memory_states, 1, 2
          )  # Transposing back
      else:
          memory_states = decoder_states
      return memory_states
  
file_path: NeMo/nemo/collections/nlp/modules/common/transformer/transformer_decoders.py
imports: import copy
  import torch
  import torch.nn
  from nemo.collections.common.parts import form_attention_mask
  from nemo.collections.nlp.modules.common.transformer.transformer_modules import MultiHeadAttention, PositionWiseFF
label_desc: Use image node_id 2 for calling the TransformerDecoder obj's underlying member method code with example usage: obj._get_memory_states(decoder_states, decoder_mems_list, i) and returns: memory_states