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: Ensures that :class:`Pipe` caches CUDA streams for copy.
  
  It's worth to cache CUDA streams although PyTorch already manages a
  pool of pre-allocated CUDA streams, because it may reduce GPU memory
  fragmentation when the number of micro-batches is small.
docstr_tok: ["Ensures","that",":","class",":","`","Pipe","`","caches","CUDA","streams","for","copy",".","It","'s","worth","to","cache","CUDA","streams","although","PyTorch","already","manages","a","pool","of","pre-allocated","CUDA","streams",",","because","it","may","reduce","GPU","memory","fragmentation","when","the","number","of","micro-batches","is","small","."]
code: def _ensure_copy_streams(self) -> List[List[AbstractStream]]:
      """Ensures that :class:`Pipe` caches CUDA streams for copy.
  
      It's worth to cache CUDA streams although PyTorch already manages a
      pool of pre-allocated CUDA streams, because it may reduce GPU memory
      fragmentation when the number of micro-batches is small.
  
      """
      if not self._copy_streams:
          for device in self.devices:
              self._copy_streams.append(
                  [new_stream(device) for _ in range(self.chunks)]
              )
  
      return self._copy_streams
  
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 8 for calling the Pipe obj's underlying member method code with example usage: obj._ensure_copy_streams() and returns: self