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: Track the rng state.
docstr_tok: ["Track","the","rng","state","."]
code: def add(self, name, seed):
      """Track the rng state."""
      # Check seed is not already used.
      if seed in self.seeds_:
          raise Exception("seed {} already exists".format(seed))
      self.seeds_.add(seed)
      # Check that state is not already defined.
      if name in self.states_:
          raise Exception(
              "cuda rng state {} already exists".format(name)
          )
      # Get the current rng state.
      orig_rng_state = torch.cuda.get_rng_state()
      # Set the new state and store it.
      torch.cuda.manual_seed(seed)
      self.states_[name] = torch.cuda.get_rng_state()
      # Reset rng state to what it was.
      _set_cuda_rng_state(orig_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 5 for calling the CudaRNGStatesTracker obj's underlying member method code with example usage: obj.add(name, seed) without return types