code: def rpc_reuse_in_final_stage():
      # 'reused' and 'reused2' are located on stage 2, so the backward pass for
      # the final stage will need to first send gradients to stage 2, then receive
      # gradients from stage 2. This tests custom logic to handle reuse of layers
      # in the final stage of the pipeline.
  
      reused = nn.Linear(10, 10)
      reused2 = nn.Linear(10, 10)
      model = [
          nn.Linear(10, 10),
          nn.ReLU(),
          nn.Linear(10, 10),
          reused2,
          nn.ReLU(),
          reused,
          nn.ReLU(),
          reused,
          reused2,
          nn.ReLU(),
          reused,
          nn.ReLU(),
      ]
      balance = [2, 3, 4]
  
      init_rpc()
  
      if torch.distributed.get_rank() != 0:
          rpc.shutdown()
          torch.distributed.barrier()
          return
  
      pipe = PipeRPCWrapper(model, balance, worker_map=get_worker_map())
  
      inputs = torch.rand(10).cuda()
      target = torch.rand(10).cuda()
      output = pipe(inputs)
      nn.MSELoss()(output, target).backward()
      output = pipe(inputs)
      nn.MSELoss()(output, target).backward()
      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 8 for calling a global function with example usage: rpc_reuse_in_final_stage() without return types