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
  from torch.cuda import device
  from torch.utils.checkpoint import detach_variable
  from megatron.core.parallel_state import get_data_parallel_rank, get_expert_model_parallel_rank, get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from megatron.core.utils import safely_set_viewless_tensor_data
  from .utils import gather_split_1d_tensor, split_tensor_into_1d_equal_chunks
label_desc: Use image node_id 6 for calling a global function with example usage: checkpoint(function, distribute_saved_activations) and returns: CheckpointFunction