original_docstring: :class:`AsyncPipe` 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 be a
          :class:`~torch.Tensor` or a tuple of tensors. This restriction is
          applied at partition boundaries too.
  
          Args:
              input (torch.Tensor or tensors): input mini-batch
  
          Returns:
              tensor or tensors: output mini-batch
  
          Raises:
              TypeError: input is not a tensor or tensors.
  
          
docstring: :class:`AsyncPipe` 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 be a
  :class:`~torch.Tensor` or a tuple of tensors. This restriction is
  applied at partition boundaries too.
docstring_tokens: : class : ` AsyncPipe ` 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 be a : class : ` ~torch . Tensor ` or a tuple of tensors . This restriction is applied at partition boundaries too .
code: def forward(self, input: TensorOrTensors, *, event=None) -> TensorOrTensors:  # type: ignore
          """:class:`AsyncPipe` 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 be a
          :class:`~torch.Tensor` or a tuple of tensors. This restriction is
          applied at partition boundaries too.
  
          Args:
              input (torch.Tensor or tensors): input mini-batch
  
          Returns:
              tensor or tensors: output mini-batch
  
          Raises:
              TypeError: input is not a tensor or tensors.
  
          """
          microbatch.check(input)
  
          if not self.pipeline:
              # No pipeline is not illegal, more ranks than partitions
              return input
  
          # Divide a mini-batch into micro-batches.
          batches = microbatch.scatter(input, self.chunks)
  
          # Run pipeline parallelism.
          with self.lock:
              self.pipeline.run(self.training, batches, event)
  
              if self.final_stage:
                  output = microbatch.gather(batches)
              else:
                  # Don't merge micro-batches to avoid unnecessary edges in autograd
                  # graph
                  # FIXME(tom) should figure out a proper type here
                  output = batches  # type: ignore
  
              return output
code_tokens: def forward ( self , input : TensorOrTensors , * , event = None ) -> TensorOrTensors : microbatch . check ( input ) if not self . pipeline : return input batches = microbatch . scatter ( input , self . chunks ) with self . lock : self . pipeline . run ( self . training , batches , event ) if self . final_stage : output = microbatch . gather ( batches ) else : output = batches return output
short_docstring: :class:`AsyncPipe` is a fairly transparent module wrapper.
short_docstring_tokens: : class : ` AsyncPipe ` is a fairly transparent module wrapper .
comment: # type: ignore$SEP$""":class:`AsyncPipe` 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 be a
          :class:`~torch.Tensor` or a tuple of tensors. This restriction is
          applied at partition boundaries too.
  
          Args:
              input (torch.Tensor or tensors): input mini-batch
  
          Returns:
              tensor or tensors: output mini-batch
  
          Raises:
              TypeError: input is not a tensor or tensors.
  
          """$SEP$# No pipeline is not illegal, more ranks than partitions$SEP$# Divide a mini-batch into micro-batches.$SEP$# Run pipeline parallelism.$SEP$# Don't merge micro-batches to avoid unnecessary edges in autograd$SEP$# graph$SEP$# FIXME(tom) should figure out a proper type here$SEP$# type: ignore