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, *args):  # type: ignore
      ctx.run_function = run_function
  
      # 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()
      )
  
      ctx.save_for_backward(*args)
      with torch.no_grad():
          outputs = run_function(*args)
      return outputs
  
imports: import contextlib
  from typing import Dict, Iterator, Set, Union
  import torch
  from torch.cuda import _lazy_call
  from torch.utils.checkpoint import detach_variable
  from .initialize import get_data_parallel_rank, get_model_parallel_rank
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) and returns: outputs