class_docstr: Tracker for the cuda RNG states.
  
  Using the `add` method, a cuda rng state is initialized based on
  the input `seed` and is assigned to `name`. Later, by forking the
  rng state, we can perform operations and return to our starting
  cuda state.
class_docstr_tok: ["Tracker","for","the","cuda","RNG","states",".","Using","the","`","add","`","method",",","a","cuda","rng","state","is","initialized","based","on","the","input","`","seed","`","and","is","assigned","to","`","name","`",".","Later",",","by","forking","the","rng","state",",","we","can","perform","operations","and","return","to","our","starting","cuda","state","."]
docstr: Fork the cuda rng state, perform operations, and exit with
  the original state.
docstr_tok: ["Fork","the","cuda","rng","state",",","perform","operations",",","and","exit","with","the","original","state","."]
code: def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
      """Fork the cuda rng state, perform operations, and exit with
      the original state."""
      # Check if we have added the state
      if name not in self.states_:
          raise Exception("cuda rng state {} is not added".format(name))
      # Store current rng state.
      orig_cuda_rng_state = torch.cuda.get_rng_state()
      # Set rng state to the desired one
      _set_cuda_rng_state(self.states_[name])
      # Do the stuff we wanted to do.
      try:
          yield
      finally:
          # Update the current rng state for later use.
          self.states_[name] = torch.cuda.get_rng_state()
          # And set the state to the original state we started with.
          _set_cuda_rng_state(orig_cuda_rng_state)
  
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 the CudaRNGStatesTracker obj's underlying member method code with example usage: obj.fork(name) without return types