code: def setup_method(self, method):
      Utils.initialize_model_parallel(1, 1)
      model_parallel_cuda_manual_seed(123)
      print("done intializing")
      num_moe_experts = 2
      transformer_config = TransformerConfig(
          num_layers=2,
          hidden_size=12,
          num_attention_heads=4,
          num_moe_experts=num_moe_experts,
          use_cpu_initialization=True,
      )
      transformer_layer_spec = (
          get_gpt_layer_with_transformer_engine_spec(
              num_experts=num_moe_experts, moe_grouped_gemm=False
          )
      )
      self.switch_mlp = SwitchMLP(
          transformer_config,
          transformer_layer_spec.submodules.mlp.submodules,
      )
  
file_path: megatron-lm/tests/unit_tests/transformer/moe/test_switch_mlp.py
imports: import pytest
  import torch
  from megatron.core.transformer.moe.switch_mlp import SwitchMLP
  from tests.unit_tests.test_utilities import Utils
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
label_desc: Use image node_id 1 for calling the TestParallelSwitchMLP obj's underlying member method code with example usage: obj.setup_method(method) without return types