code: def create_task(
      checkpoint_stop: int,
      chunk_id: int,
      part_id: int,
      batch: Batch,
      partition: nn.Sequential,
      skip_trackers: List[SkipTrackerThroughPotals],
  ) -> Task:
      # Determine whether checkpointing or not.
      if chunk_id < checkpoint_stop:
  
          def function(
              input: TensorOrTensors,
              partition: nn.Sequential = partition,
              skip_tracker: SkipTrackerThroughPotals = skip_trackers[
                  chunk_id
              ],
              chunk_id: int = chunk_id,
              part_id: int = part_id,
          ) -> TensorOrTensors:
              with use_skip_tracker(skip_tracker), record_function(
                  "chunk%d-part%d" % (chunk_id, part_id)
              ):
                  ret = partition(input)
                  # We do a check here because the backtrace from the checkpoint backward code path
                  # is very hard to make sense. It would be much easier to check earlier at this point.
                  assert (
                      type(ret) is not list
                  ), "Only Tensor or Tuple of Tensor output is supported"
                  return ret
  
          chk = Checkpointing(function, batch)
          task = Task(
              None, compute=chk.checkpoint, finalize=chk.recompute
          )
          del function, chk  # TODO(tom) maybe remove
  
      else:
  
          def compute(
              batch: Batch = batch,
              partition: nn.Sequential = partition,
              skip_tracker: SkipTrackerThroughPotals = skip_trackers[
                  chunk_id
              ],
              chunk_id: int = chunk_id,
              part_id: int = part_id,
          ) -> Batch:
              with use_skip_tracker(skip_tracker), record_function(
                  "chunk%d-part%d" % (chunk_id, part_id)
              ):
                  return batch.call(partition)
  
          task = Task(None, compute=compute, finalize=None)
          del compute  # TODO(tom) maybe remove
  
      return task
  
imports: from collections import OrderedDict
  from dataclasses import dataclass
  from enum import Enum, auto
  from threading import Event
  from typing import Dict, Iterable, List, Optional, Tuple
  import torch
  from torch import Tensor, nn
  from torch.autograd.profiler import record_function
  from torch.distributed import ProcessGroup
  from fairscale.nn.model_parallel import get_pipeline_parallel_ranks
  from .checkpoint import Checkpointing
  from .messages import Transport
  from .microbatch import Batch
  from .skip.tracker import SkipTrackerThroughPotals, use_skip_tracker
  from .types import EVENT_LOOP_QUEUE, PipeMessage, TensorOrTensors, Tensors
  from .worker import Task
label_desc: Use image node_id 1 for calling a global function with example usage: create_task(checkpoint_stop, chunk_id, part_id, batch, partition, skip_trackers) and returns: task, ret, batch