code: def checkpoint(
      function, activation_offload, *args, use_reentrant: bool = True
  ):
      """Checkpoint the computation while preserve the rng states, modified from Pytorch torch.utils.checkpoint.
  
      Args:
          function: Describe the forward pass function. It should know how to handle the input tuples.
          activation_offload: The variable to check whether we should offload activation to cpu
          args (list): Tuple containing the parameters of the function
          use_reentrant: Bool type to check if we need to use_reentrant, if use_reentrant=False, there
          might be more flexibility for user to define there checkpoint function
  
      Returns:
          Output of running function with provided args.
      """
      if use_reentrant:
          return CheckpointFunction.apply(
              function, activation_offload, *args
          )
      else:
          return _checkpoint_without_reentrant(
              function,
              activation_offload,
              *args,
          )
  
file_path: ColossalAI/colossalai/legacy/utils/activation_checkpoint.py
imports: import weakref
  import torch
  from torch.utils.checkpoint import check_backward_validity, detach_variable
  from colossalai.legacy.context.random import get_current_mode, get_states, set_mode, set_seed_states, sync_states
  from colossalai.utils.device import autocast, get_current_device
label_desc: Use image node_id 2 for calling a global function with example usage: checkpoint(function, activation_offload) and returns: CheckpointFunction, _checkpoint_without_reentrant