code: def get_modules(cls, config, use_transformer_engine, use_gpu):
      # Retro decoder layer.
      decoder_block_spec = get_retro_decoder_block_spec(
          config, use_transformer_engine=use_transformer_engine
      )
      decoder_block = TransformerBlock(
          config=config, spec=decoder_block_spec
      )
      decoder_layers = [
          layer
          for layer in decoder_block.layers
          if isinstance(
              layer.cross_attention, RetroDecoderCrossAttention
          )
      ]
      decoder_layer = decoder_layers[0]
  
      # Retro encoder layer.
      encoder_block = decoder_layer.cross_attention.encoder
      encoder_layers = [
          layer
          for layer in encoder_block.layers
          if isinstance(
              layer.cross_attention, RetroEncoderCrossAttention
          )
      ]
      encoder_layer = encoder_layers[0]
  
      # Modules.
      modules = types.SimpleNamespace(
          decoder_attn=decoder_layer.cross_attention,
          decoder_bda=decoder_layer.cross_attn_bda,
          encoder_attn=encoder_layer.cross_attention,
          encoder_bda=encoder_layer.cross_attn_bda,
          encoder_norm=encoder_layer.pre_mlp_layernorm,
      )
  
      # GPU.
      if use_gpu:
          [m.cuda() for m in vars(modules).values()]
  
      return modules
  
file_path: megatron-lm/tests/unit_tests/transformer/test_retro_attention.py
imports: import torch
  import types
  from megatron.core.models.retro import RetroConfig, get_retro_decoder_block_spec
  from megatron.core.models.retro.decoder_attention import RetroDecoderCrossAttention, RetroDecoderBiasDropoutAdd
  from megatron.core.models.retro.encoder_attention import RetroEncoderCrossAttention, RetroEncoderBiasDropoutAdd, RetroEncoderLayerNorm
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.transformer.transformer_block import TransformerBlock
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 2 for calling the TestRetroAttention obj's underlying member method code with example usage: obj.get_modules(cls, config, use_transformer_engine, use_gpu) and returns: modules