code: def check_pipe_against_reference(
      balance,
      model_constructor,
      checkpoint="except_last",
      custom_inputs=None,
  ):
      model = model_constructor()
      reference_model = model_constructor()
      for src, dst in zip(model, reference_model):
          dst.load_state_dict(copy.deepcopy(src.state_dict()))
  
      reference_model = nn.Sequential(*reference_model).cuda()
  
      pipe = PipeRPCWrapper(
          model,
          balance,
          input_device=torch.cuda.current_device(),
          worker_map=get_worker_map(),
          checkpoint=checkpoint,
      )
  
      pipe.foreach_worker(register_optimizer, include_self=True)
      register_optimizer(None, reference_model)
  
      inputs = torch.rand(10).cuda()
      target = torch.rand(10).cuda()
      cloned = inputs.clone()
      output = pipe(inputs)
      ref_out = reference_model(inputs)
  
      assert torch.equal(ref_out.cpu(), output.cpu())
  
      for out in output, ref_out:
          target = target.to(out.device)
          loss = nn.MSELoss()(out, target)
          loss.backward()
  
      pipe.foreach_worker(step_optimizer, include_self=True)
      step_optimizer(None, reference_model.cuda())
  
      pipe.eval()
      reference_model.eval()
  
      final_output = pipe(inputs)
      final_ref = reference_model(inputs.cuda())
  
      assert torch.equal(final_output.cpu(), final_ref.cpu())
  
file_path: fairscale/tests/nn/pipe_process/test_rpc.py
imports: import copy
  import os
  import pytest
  import torch
  from torch import nn
  from torch.distributed import rpc
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.internal import torch_version
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from fairscale.nn.pipe import PipeRPCWrapper
label_desc: Use image node_id 5 for calling a global function with example usage: check_pipe_against_reference(balance, model_constructor, checkpoint, custom_inputs) without return types