class_docstr: Wraps an arbitrary :class:`nn.Sequential <torch.nn.Sequential>` module
  to train on Pipe_. If the module requires lots of memory, Pipe will be
  very efficient.
  
  Pipe combines pipeline parallelism with checkpointing to reduce peak
  memory required to train while minimizing device under-utilization.
  
  You should determine the balance when defining a :class:`AsyncPipe` module, as
  balancing will not be done automatically. The module will be partitioned
  into multiple devices according to the given balance. You may rely on
  heuristics to find your own optimal configuration.
  
  Args:
      module (torch.nn.Sequential):
          sequential module to be parallelized
      balance (ints):
          list of number of layers in each partition
  
  Keyword Args:
      group (ProcessGroup):
          the process group that all
          pipeline stages are a member of. Defaults to
          `get_pipeline_parallel_group()`
      worker_map (Dict[int, str]):
          a map from worker name (the first argument to
          `torch.distributed.rpc.init_rpc`) to global rank (i.e.
          `torch.distributed.get_rank()`) needed in order for pipeline stages
          to communicate with each other
      input_device (device):
          the device on which tensors should be located before being passed to
          the first module in a given pipeline stage
      chunks (int):
          number of micro-batches (default: ``1``)
      checkpoint (str):
          when to enable checkpointing, one of ``'always'``,
          ``'except_last'``, or ``'never'`` (default: ``'except_last'``)
      deferred_batch_norm (bool):
          whether to use deferred BatchNorm moving statistics (default:
          :data:`False`, see :class:`DeferredBatchNorm` for more
          details)
  
  Raises:
      TypeError:
          the module is not a :class:`nn.Sequential <torch.nn.Sequential>`.
      ValueError:
          invalid arguments, or wrong balance
      IndexError:
          the number of devices is fewer than the number of partitions.
class_docstr_tok: ["Wraps","an","arbitrary",":","class",":","`","nn.Sequential","<","torch.nn.Sequential",">","`","module","to","train","on","Pipe_",".","If","the","module","requires","lots","of","memory",",","Pipe","will","be","very","efficient",".","Pipe","combines","pipeline","parallelism","with","checkpointing","to","reduce","peak","memory","required","to","train","while","minimizing","device","under-utilization",".","You","should","determine","the","balance","when","defining","a",":","class",":","`","AsyncPipe","`","module",",","as","balancing","will","not","be","done","automatically",".","The","module","will","be","partitioned","into","multiple","devices","according","to","the","given","balance",".","You","may","rely","on","heuristics","to","find","your","own","optimal","configuration",".","Args",":","module","(","torch.nn.Sequential",")",":","sequential","module","to","be","parallelized","balance","(","ints",")",":","list","of","number","of","layers","in","each","partition","Keyword","Args",":","group","(","ProcessGroup",")",":","the","process","group","that","all","pipeline","stages","are","a","member","of",".","Defaults","to","`","get_pipeline_parallel_group","(",")","`","worker_map","(","Dict","[","int",",","str","]",")",":","a","map","from","worker","name","(","the","first","argument","to","`","torch.distributed.rpc.init_rpc","`",")","to","global","rank","(","i.e",".","`","torch.distributed.get_rank","(",")","`",")","needed","in","order","for","pipeline","stages","to","communicate","with","each","other","input_device","(","device",")",":","the","device","on","which","tensors","should","be","located","before","being","passed","to","the","first","module","in","a","given","pipeline","stage","chunks","(","int",")",":","number","of","micro-batches","(","default",":","``","1","``",")","checkpoint","(","str",")",":","when","to","enable","checkpointing",",","one","of","``","'always","'","``",",","``","'except_last","'","``",",","or","``","'never","'","``","(","default",":","``","'except_last","'","``",")","deferred_batch_norm","(","bool",")",":","whether","to","use","deferred","BatchNorm","moving","statistics","(","default",":",":","data",":","`","False","`",",","see",":","class",":","`","DeferredBatchNorm","`","for","more","details",")","Raises",":","TypeError",":","the","module","is","not","a",":","class",":","`","nn.Sequential","<","torch.nn.Sequential",">","`",".","ValueError",":","invalid","arguments",",","or","wrong","balance","IndexError",":","the","number","of","devices","is","fewer","than","the","number","of","partitions","."]
docstr: Gets a layer in the underlying sequential module.
docstr_tok: ["Gets","a","layer","in","the","underlying","sequential","module","."]
code: def __getitem__(self, index: int) -> nn.Module:
      """Gets a layer in the underlying sequential module."""
      partitions: List[Any]
      partitions = self.partitions
  
      if index < 0:
          partitions = partitions[::-1]
  
      for partition in partitions:
          try:
              if isinstance(partition, ModuleWrapper):
                  return partition.module[index]
              else:
                  return partition[index]
          except IndexError:
              pass
  
          shift = len(partition)
  
          if index < 0:
              index += shift
          else:
              index -= shift
  
      raise IndexError
  
imports: from collections import OrderedDict
  from dataclasses import dataclass, field
  import itertools
  import threading
  from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
  import warnings
  import torch
  from torch import Tensor, nn
  from fairscale.nn.model_parallel import get_pipeline_parallel_group
  from .None import microbatch
  from .async_pipeline import AsyncPipeline
  from .async_schedule import Invocation, Location, ModuleWrapper
  from .batchnorm import DeferredBatchNorm
  from .skip.layout import SkipLayout
  from .skip.skippable import Skippable
  from .types import LazyModule
label_desc: Use image node_id 5 for calling the AsyncPipe obj's underlying member method code with example usage: obj.__getitem__(index) and returns: partition, partition