code: def checkpoint(function, distribute_saved_activations, *args):
      """Checkpoint a model or part of the model.
      This has been directly copied from torch.utils.checkpoint."""
      return CheckpointFunction.apply(
          function, distribute_saved_activations, *args
      )
  
imports: import contextlib
  import torch
  from torch import _C
  from torch.cuda import _lazy_call, device
  from torch.utils.checkpoint import detach_variable
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.tensor_parallel.memory import allocate_mem_buff
  from apex.transformer.utils import split_tensor_into_1d_equal_chunks
  from apex.transformer.utils import gather_split_1d_tensor
label_desc: Use image node_id 6 for calling a global function with example usage: checkpoint(function, distribute_saved_activations) and returns: CheckpointFunction