class_docstr: This function is adapted from torch.utils.checkpoint with
  two main changes:
      1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
      2) the states in the model parallel tracker are also properly
         tracked/set/reset.
class_docstr_tok: ["This","function","is","adapted","from","torch.utils.checkpoint","with","two","main","changes",":","1",")","torch.cuda.set_rng_state","is","replaced","with","`","_set_cuda_rng_state","`","2",")","the","states","in","the","model","parallel","tracker","are","also","properly","tracked\/set\/reset","."]
code: def forward(ctx, run_function, distribute_saved_activations, *args):
      ctx.run_function = run_function
      ctx.distribute_saved_activations = distribute_saved_activations
  
      # Copy the rng states.
      ctx.fwd_cpu_rng_state = torch.get_rng_state()
      ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
      ctx.fwd_cuda_rng_state_tracker = (
          get_cuda_rng_tracker().get_states()
      )
  
      with torch.no_grad():
          outputs = run_function(*args)
  
      # Divide hidden states across model parallel group and only keep
      # the chunk corresponding to the current rank.
      if ctx.distribute_saved_activations:
          ctx.input_0_shape = args[0].shape
  
      # Store everything.
      ctx.save_for_backward(*args)
      return outputs
  
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 1 for calling the CheckpointFunction obj's underlying member method code with example usage: obj.forward(ctx, run_function, distribute_saved_activations) and returns: outputs