code: def test(self, ep_size, use_residual):
      if not required_torch_version():
          pytest.skip(
              "DeepSpeed MoE tests need torch 1.8 or higher to run correctly"
          )
  
      config_dict = {
          "train_batch_size": 8,
          "steps_per_print": 1,
          "fp16": {"enabled": True},
      }
      hidden_dim = 16
  
      # E+D -- ep_size = 2
      # E only -- ep_size = 4
      model = SimplePRMoEModel(
          hidden_dim, ep_size=ep_size, use_residual=use_residual
      )
      optimizer = torch.optim.AdamW(params=model.parameters())
      model, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          optimizer=optimizer,
          dist_init_required=False,
      )
  
      data_loader = sequence_dataloader(
          model=model,
          total_samples=50,
          hidden_dim=hidden_dim,
          device=model.device,
      )
  
      for n, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          model.backward(loss)
          model.step()
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/moe/test_moe.py
imports: import torch
  import deepspeed
  import pytest
  from unit.common import DistributedTest
  from unit.simple_model import SimplePRMoEModel, SimpleMoEModel, sequence_dataloader
  from unit.util import required_torch_version
label_desc: Use image node_id 1 for calling the TestPRMoE obj's underlying member method code with example usage: obj.test(ep_size, use_residual) without return types