class_docstr: A wrapper for Pipe to control the entire pipeline from a single process.
  Typical usecase would have rank 0 construct `PipeRPCWrapper` and run the
  training loop as normal, and all other ranks would call
  `torch.distributed.rpc.shutdown()`
  
  To run code on each worker, e.g. to run the optimizer, use `foreach_worker`
class_docstr_tok: ["A","wrapper","for","Pipe","to","control","the","entire","pipeline","from","a","single","process",".","Typical","usecase","would","have","rank","0","construct","`","PipeRPCWrapper","`","and","run","the","training","loop","as","normal",",","and","all","other","ranks","would","call","`","torch.distributed.rpc.shutdown","(",")","`","To","run","code","on","each","worker",",","e.g",".","to","run","the","optimizer",",","use","`","foreach_worker","`"]
code: def forward(self, tensor: TensorOrTensors) -> TensorOrTensors:  # type: ignore
      shape = get_shapes(tensor)
      dtype = get_dtype(tensor)
  
      if isinstance(tensor, torch.Tensor):
          num_tensors = 1
      else:
          num_tensors = len(tensor)
  
      futures = [
          rpc.rpc_async(
              self._get_rpc_name(rank),
              self._model_forward,
              args=(self.model.training, shape, dtype),
          )
          for rank in range(1, self.group.size())
      ]
  
      if self.model.final_stage:
          return self.model(tensor)
      else:
          event = Event()
          t = Thread(
              target=self._model_forward_first_stage,
              args=(tensor, event),
          )
          t.start()
  
          shape, dtype = futures.pop().wait()
          dest_rank = self.group.size() - 1
          dest = self._get_rpc_name(dest_rank)
          dest_global_rank = _get_global_rank(self.group, dest_rank)
          src_global_rank = torch.distributed.get_rank()
          queue = EVENT_LOOP_QUEUE
  
          activations = PipeMessage(
              dest_global_rank,
              src_global_rank,
              queue_name=queue,
              tensor_count=num_tensors,
          )
          grads = PipeMessage(
              src_global_rank,
              dest_global_rank,
              queue_name=queue,
              tensor_count=num_tensors,
          )
  
          back_fut = rpc.rpc_async(
              dest,
              self._send_result_and_do_backwards,
              args=(self.model.training, activations, grads),
          )
          futures.append(back_fut)
  
          result = self._recv_result(
              self.model, shape, dtype, activations
          )
          if isinstance(result, torch.Tensor):
              result.requires_grad_()
          else:
              for r in result:
                  r.requires_grad_()
  
          assert self.model.pipeline
          return PipeBackRedirect.apply(
              result,
              dest_global_rank,
              event,
              grads,
              self.model.pipeline.transport,
              futures,
          )
  
imports: from threading import Event, Lock, Thread
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
  import torch
  from torch import nn
  from torch.distributed import ProcessGroup, rpc
  from torch.distributed.distributed_c10d import _get_global_rank
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from .async_pipe import AsyncPipe
  from .types import EVENT_LOOP_QUEUE, PipeMessage, TensorOrTensors
label_desc: Use image node_id 5 for calling the PipeRPCWrapper obj's underlying member method code with example usage: obj.forward(tensor) and returns: self, PipeBackRedirect