class_docstr: Wraps an arbitrary :class:`nn.Sequential <torch.nn.Sequential>` module
  to train on using synchronous pipeline parallelism. If the module requires
  lots of memory and doesn't fit on a single GPU, pipeline parallelism is a
  useful technique to employ for training.
  
  The implementation is based on the torchgpipe_ paper.
  
  .. _torchgpipe: https://arxiv.org/abs/2004.09910
  
  Pipe combines pipeline parallelism with checkpointing to reduce peak
  memory required to train while minimizing device under-utilization.
  
  You should place all the modules on the appropriate devices and wrap them
  into an :class:`nn.Sequential <torch.nn.Sequential>` module defining the
  desired order of execution. If a module does not contain any
  parameters/buffers, it is assumed this module should be executed on CPU
  and appropriate input tensors to the module are moved to CPU before
  execution. This behavior can be overridden by the :class:`WithDevice`
  wrapper which can be used to explicitly specify which device a module
  should run on.
  
  Args:
      module (:class:`nn.Sequential <torch.nn.Sequential>`):
          sequential module to be parallelized using pipelining. Each module
          in the sequence has to have all of its parameters on a single
          device. Each module in the sequence has to either be an nn.Module
          or :class:`nn.Sequential <torch.nn.Sequential>` (to combine multiple
          sequential modules on a single device)
      chunks (int):
          number of micro-batches (default: ``1``)
      checkpoint (str):
          when to enable checkpointing, one of ``'always'``,
          ``'except_last'``, or ``'never'`` (default: ``'except_last'``).
          ``'never'`` disables checkpointing completely, ``'except_last'``
          enables checkpointing for all micro-batches except the last one
          and ``'always'`` enables checkpointing for all micro-batches.
      deferred_batch_norm (bool):
          whether to use deferred ``BatchNorm`` moving statistics (default:
          :data:`False`). If set to :data:`True`, we track statistics across
          multiple micro-batches to update the running statistics per
          mini-batch.
  
  Raises:
      TypeError:
          the module is not a :class:`nn.Sequential <torch.nn.Sequential>`.
      ValueError:
          invalid arguments
  
  Example::
      Pipeline of two FC layers across GPUs 0 and 1.
  
      >>> # Need to initialize RPC framework first.
      >>> # xdoctest: +SKIP
      >>> os.environ['MASTER_ADDR'] = 'localhost'
      >>> os.environ['MASTER_PORT'] = '29500'
      >>> torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1)
      >>>
      >>> # Build pipe.
      >>> fc1 = nn.Linear(16, 8).cuda(0)
      >>> fc2 = nn.Linear(8, 4).cuda(1)
      >>> model = nn.Sequential(fc1, fc2)
      >>> model = Pipe(model, chunks=8)
      >>> input = torch.rand(16, 16).cuda(0)
      >>> output_rref = model(input)
  
  .. note::
      You can wrap a :class:`Pipe` model with
      :class:`torch.nn.parallel.DistributedDataParallel` only when the
      checkpoint parameter of :class:`Pipe` is ``'never'``.
  
  .. note::
      :class:`Pipe` only supports intra-node pipelining currently, but
      will be expanded to support inter-node pipelining in the future.
      The forward function returns an :class:`~torch.distributed.rpc.RRef`
      to allow for inter-node pipelining in the future, where the output
      might be on a remote host. For intra-node pipelining you can use
      :meth:`~torch.distributed.rpc.RRef.local_value` to retrieve the
      output locally.
  
  .. warning::
      :class:`Pipe` is experimental and subject to change.
docstr: Processes a single input mini-batch through the pipe and returns an
  :class:`~torch.distributed.rpc.RRef` pointing to the output.
  :class:`Pipe` is a fairly transparent module wrapper. It doesn't
  modify the input and output signature of the underlying module. But
  there's type restriction. Input and output have to contain at least one
  tensor. This restriction is applied at partition boundaries too.
  
  The sequence of inputs are fed into the first stage of the pipeline as
  ``*inputs``. As a result the positional args for this function should
  match the positional args for the first stage of the pipeline. The same
  condition applies for output of one stage of the pipeline which is the
  input for the next stage.
  
  The input tensor is split into multiple micro-batches based on the
  ``chunks`` parameter used to initialize :class:`Pipe`. The batch size
  is assumed to be the first dimension of the tensor and if the batch
  size is less than ``chunks``, the number of micro-batches is equal to
  the batch size.
  
  Only tensors are split into multiple micro-batches, non-Tensor inputs
  are just replicated as-is in each micro-batch. For non-Tensor outputs
  in the last stage of the pipeline, they are aggregated as a ``List``
  and returned the user. For example, if you have 2 micro-batches
  returning the integer 5, the user would receive the consolidated
  output of `[5, 5]`
  
  All the input tensors need to be on the same device as the first
  partition of the pipeline.
  
  If a tensor is wrapped with the :class:`NoChunk` wrapper, the tensor
  is not split across micro-batches and is replicated as-is similar to
  non-tensors.
  
  Args:
      inputs: input mini-batch
  
  Returns:
      :class:`~torch.distributed.rpc.RRef` to the output of the mini-batch
  
  Raises:
      TypeError: input doesn't contain at least one tensor
docstr_tok: ["Processes","a","single","input","mini-batch","through","the","pipe","and","returns","an",":","class",":","`","~torch.distributed.rpc.RRef","`","pointing","to","the","output",".",":","class",":","`","Pipe","`","is","a","fairly","transparent","module","wrapper",".","It","doesn't","modify","the","input","and","output","signature","of","the","underlying","module",".","But","there","'s","type","restriction",".","Input","and","output","have","to","contain","at","least","one","tensor",".","This","restriction","is","applied","at","partition","boundaries","too",".","The","sequence","of","inputs","are","fed","into","the","first","stage","of","the","pipeline","as","``","*","inputs","``",".","As","a","result","the","positional","args","for","this","function","should","match","the","positional","args","for","the","first","stage","of","the","pipeline",".","The","same","condition","applies","for","output","of","one","stage","of","the","pipeline","which","is","the","input","for","the","next","stage",".","The","input","tensor","is","split","into","multiple","micro-batches","based","on","the","``","chunks","``","parameter","used","to","initialize",":","class",":","`","Pipe","`",".","The","batch","size","is","assumed","to","be","the","first","dimension","of","the","tensor","and","if","the","batch","size","is","less","than","``","chunks","``",",","the","number","of","micro-batches","is","equal","to","the","batch","size",".","Only","tensors","are","split","into","multiple","micro-batches",",","non-Tensor","inputs","are","just","replicated","as-is","in","each","micro-batch",".","For","non-Tensor","outputs","in","the","last","stage","of","the","pipeline",",","they","are","aggregated","as","a","``","List","``","and","returned","the","user",".","For","example",",","if","you","have","2","micro-batches","returning","the","integer","5",",","the","user","would","receive","the","consolidated","output","of","`","[","5",",","5","]","`","All","the","input","tensors","need","to","be","on","the","same","device","as","the","first","partition","of","the","pipeline",".","If","a","tensor","is","wrapped","with","the",":","class",":","`","NoChunk","`","wrapper",",","the","tensor","is","not","split","across","micro-batches","and","is","replicated","as-is","similar","to","non-tensors",".","Args",":","inputs",":","input","mini-batch","Returns",":",":","class",":","`","~torch.distributed.rpc.RRef","`","to","the","output","of","the","mini-batch","Raises",":","TypeError",":","input","does","n't","contain","at","least","one","tensor"]
code: def forward(self, *inputs) -> RRef:
      """
      Processes a single input mini-batch through the pipe and returns an
      :class:`~torch.distributed.rpc.RRef` pointing to the output.
      :class:`Pipe` is a fairly transparent module wrapper. It doesn't
      modify the input and output signature of the underlying module. But
      there's type restriction. Input and output have to contain at least one
      tensor. This restriction is applied at partition boundaries too.
  
      The sequence of inputs are fed into the first stage of the pipeline as
      ``*inputs``. As a result the positional args for this function should
      match the positional args for the first stage of the pipeline. The same
      condition applies for output of one stage of the pipeline which is the
      input for the next stage.
  
      The input tensor is split into multiple micro-batches based on the
      ``chunks`` parameter used to initialize :class:`Pipe`. The batch size
      is assumed to be the first dimension of the tensor and if the batch
      size is less than ``chunks``, the number of micro-batches is equal to
      the batch size.
  
      Only tensors are split into multiple micro-batches, non-Tensor inputs
      are just replicated as-is in each micro-batch. For non-Tensor outputs
      in the last stage of the pipeline, they are aggregated as a ``List``
      and returned the user. For example, if you have 2 micro-batches
      returning the integer 5, the user would receive the consolidated
      output of `[5, 5]`
  
      All the input tensors need to be on the same device as the first
      partition of the pipeline.
  
      If a tensor is wrapped with the :class:`NoChunk` wrapper, the tensor
      is not split across micro-batches and is replicated as-is similar to
      non-tensors.
  
      Args:
          inputs: input mini-batch
  
      Returns:
          :class:`~torch.distributed.rpc.RRef` to the output of the mini-batch
  
      Raises:
          TypeError: input doesn't contain at least one tensor
  
      """
      first_partition_device = (
          self.devices[0]
          if len(self.devices) != 0
          else torch.device("cpu")
      )
      microbatch.check(first_partition_device, *inputs)
  
      if not self.devices:
          # Empty sequential module is not illegal.
          return RRef(*inputs)
  
      # Divide a mini-batch into micro-batches.
      batches = microbatch.scatter(*inputs, chunks=self.chunks)
  
      # Run pipeline parallelism.
      self.pipeline.run(batches)
  
      # Merge the micro-batches into one mini-batch.
      output = microbatch.gather(batches)
      return RRef(output)
  
imports: from collections import OrderedDict
  from typing import TYPE_CHECKING, Any, Iterable, Iterator, List, Optional, Union, Sequence, Tuple, cast
  import torch
  from torch import Tensor, nn
  from torch.distributed.rpc import RRef
  import torch.autograd
  import torch.cuda
  from .None import microbatch
  from .batchnorm import DeferredBatchNorm
  from .pipeline import Pipeline
  from .skip.layout import inspect_skip_layout
  from .skip.skippable import verify_skippables
  from .stream import AbstractStream, new_stream
label_desc: Use image node_id 9 for calling the Pipe obj's underlying member method code with example usage: obj.forward() and returns: RRef, RRef