class_docstr: This class processes a single partition of the pipeline.
  Args:
      module_rref: RRef to the nn.Module for this partition. It should be on the local rpc worker.
      device: The device that holds the module.
      num_inputs: Numer of inputs to the module
      num_outputs: Number of outputs of the module. If the module output is not a tuple (and it is a
          single tensor), num_outputs should be None.
      rank: The rank of the partition
      chunks: Number of micor-batches in a mini-batch
      checkpoint_stop:: Checkpointing is done only for the first checkpoint_stop chunks of a mini-batch.
class_docstr_tok: ["This","class","processes","a","single","partition","of","the","pipeline",".","Args",":","module_rref",":","RRef","to","the","nn.Module","for","this","partition",".","It","should","be","on","the","local","rpc","worker",".","device",":","The","device","that","holds","the","module",".","num_inputs",":","Numer","of","inputs","to","the","module","num_outputs",":","Number","of","outputs","of","the","module",".","If","the","module","output","is","not","a","tuple","(","and","it","is","a","single","tensor",")",",","num_outputs","should","be","None",".","rank",":","The","rank","of","the","partition","chunks",":","Number","of","micor-batches","in","a","mini-batch","checkpoint_stop",":",":","Checkpointing","is","done","only","for","the","first","checkpoint_stop","chunks","of","a","mini-batch","."]
docstr: Runs tasks with synchronization to tensor-pipe streams.
docstr_tok: ["Runs","tasks","with","synchronization","to","tensor-pipe","streams","."]
code: def compute(
      self, pipeline_record: DistributedPipelineRecord, chunk: int
  ) -> None:
      """Runs tasks with synchronization to tensor-pipe streams."""
      checkpoint_stop = self.checkpoint_stop
  
      # Disable checkpointing if in eval mode.
      if not self.module.training:
          checkpoint_stop = 0
  
      exc_info: Optional[ExcInfo] = None
  
      batch = pipeline_record.get_batch(chunk)
  
      if is_cuda(self.stream):
          pipeline_record.sync_stream(chunk, as_cuda(self.stream))
  
      # Determine whether checkpointing or not.
      checkpoint = chunk < checkpoint_stop
      if checkpoint:
  
          def function(
              input: TensorOrTensors, chunk_id: int = chunk
          ) -> TensorOrTensors:
              with record_function(
                  "chunk%d-rank%d" % (chunk_id, pipeline_record.rank)
              ):
                  result = self.module(*input)
                  if self.num_outputs is None:
                      result = (result,)
                  return tuple(result)
  
          chk = Checkpointing(function, batch)
          task = Task(
              self.stream,
              compute=chk.checkpoint,
              finalize=chk.recompute,
          )
          del function, chk
  
      else:
  
          def compute(
              batch: Batch = batch,
              chunk_id: int = chunk,
              rank: int = pipeline_record.rank
              if pipeline_record is not None
              else -1,
          ) -> Batch:
              with record_function(
                  "chunk%d-rank%d" % (chunk_id, pipeline_record.rank)
              ):
                  result = self.module(*batch.tensors)
                  if self.num_outputs is None:
                      result = (result,)
              return Batch(result, chunk_id)
  
          task = Task(self.stream, compute=compute, finalize=None)
          del compute
  
      self.in_queue.put(task)
  
      ok, payload = self.out_queue.get()
  
      # Hold the first exception.
      if exc_info is not None:
          pass
      elif not ok:
          exc_info = cast(ExcInfo, payload)
      else:
          task, batch = cast(Tuple[Task, Batch], payload)
  
          with use_device(self.device):
              task.finalize(batch)
  
          pipeline_record.batches[chunk] = batch
  
      if exc_info is not None:
          raise exc_info[0].with_traceback(exc_info[1], exc_info[2])
  
imports: from threading import Condition
  from types import TracebackType
  from typing import Dict, List, Optional, Tuple, Type, Union, cast
  import torch
  from torch import Tensor
  from torch.autograd.profiler import record_function
  from torch.distributed import rpc
  from fairscale.nn.pipe import microbatch
  from fairscale.nn.pipe.checkpoint import Checkpointing, TensorOrTensors
  from fairscale.nn.pipe.dependency import fork, join
  from fairscale.nn.pipe.microbatch import Batch
  from fairscale.nn.pipe.stream import as_cuda, current_stream, is_cuda, use_device, use_stream
  from fairscale.nn.pipe.worker import Task, create_workers
  from .data import DataConsumer
label_desc: Use image node_id 6 for calling the PartitionHandler obj's underlying member method code with example usage: obj.compute(pipeline_record, chunk) and returns: Batch, tuple