code: def test(self, sequential_model, simple_config, batch_input):
      base_model = copy.deepcopy(sequential_model)
      base_input = batch_input.clone().detach()
      base_output = base_model(base_input)
      base_output = base_output
      base_params = sum(p.numel() for p in base_model.parameters())
  
      pipe_model = copy.deepcopy(sequential_model)
      pipe_model = PipelineModule(layers=pipe_model, num_stages=2)
  
      # Ensure all parameters are accounted for.
      my_params = sum(p.numel() for p in pipe_model.parameters())
      total_pipe_params = torch.LongTensor([my_params]).to("cuda")
      dist.all_reduce(total_pipe_params)
      total_pipe_params = total_pipe_params.item()
      assert total_pipe_params == base_params
  
      pipe_model, _, _, _ = deepspeed.initialize(
          config=simple_config,
          model=pipe_model,
          model_parameters=[p for p in pipe_model.parameters()],
      )
  
      if pipe_model.is_first_stage or pipe_model.is_last_stage:
          pipe_input = base_input.clone().detach().to("cuda")
          # label 0 is meaningless
          dataset = [(pipe_input, 0)]
          loader = RepeatingLoader(dataset)
          data_iter = iter(loader)
      else:
          data_iter = None
  
      pipe_output = pipe_model.eval_batch(data_iter=data_iter)
  
      base_output = base_output.to("cpu")
      pipe_output = pipe_output.to("cpu")
  
      assert torch.allclose(base_output, pipe_output, atol=1e-4)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/pipe/test_pipe_module.py
imports: import copy
  import torch
  import torch.nn
  import deepspeed.comm
  import pytest
  import deepspeed
  from deepspeed.pipe import PipelineModule
  from deepspeed.utils import RepeatingLoader
  from unit.common import DistributedTest
label_desc: Use image node_id 1 for calling the TestPipeModuleSequential obj's underlying member method code with example usage: obj.test(sequential_model, simple_config, batch_input) without return types