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","."]
code: def __init__(self):
      # Map from a string name to the cuda rng state.
      self.states_ = {}
      # Seeds are just for book keeping and ensure no seed is set twice.
      self.seeds_ = set()
  
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 to create a new CudaRNGStatesTracker object  with example: obj = CudaRNGStatesTracker()