code: def basic_rpc():
      init_rpc()
      if torch.distributed.get_rank() != 0:
          rpc.shutdown()
          torch.distributed.barrier()
          return
  
      model = [nn.Linear(10, 10), nn.ReLU()]
      pipe = PipeRPCWrapper(
          model,
          [1, 1],
          input_device=torch.cuda.current_device(),
          worker_map=get_worker_map(),
      )
  
      pipe.foreach_worker(register_optimizer, include_self=True)
  
      inputs = torch.rand(10).cuda()
      output = pipe(inputs)
      loss = output.mean()
      loss.backward()
  
      pipe.foreach_worker(step_optimizer, include_self=True)
  
      pipe.eval()
  
      rpc.shutdown()
      torch.distributed.barrier()
  
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 2 for calling a global function with example usage: basic_rpc() without return types