code: def test_gpu_forward(self):
      self.switch_mlp_smm.cuda()
      self.switch_mlp_gmm.cuda()
      # [sequence length, batch size, hidden size]
      seq_len = 3  # 32
      batch_size = 2
      hidden_states = torch.rand(
          (seq_len, batch_size, self.switch_mlp_smm.config.hidden_size),
          dtype=torch.bfloat16,
      )
      hidden_states = hidden_states.cuda()
      output_smm, _ = self.switch_mlp_smm(hidden_states)
      output_gmm, _ = self.switch_mlp_gmm(hidden_states)
  
file_path: megatron-lm/tests/unit_tests/transformer/moe/test_grouped_mlp.py
imports: import pytest
  import torch
  import torch.nn.functional
  from megatron.arguments import parse_args
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
  from megatron.core.transformer.moe.grouped_mlp import GroupedMLP
  from megatron.core.transformer.moe.switch_mlp import SwitchMLP
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.initialize import _set_random_seed
  from megatron.model import Float16Module
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 5 for calling the TestParallelGroupedMLP obj's underlying member method code with example usage: obj.test_gpu_forward() without return types