code: def _set_cuda_rng_state(new_state, device=-1):
      """Sets the random number generator state of the current GPU.
  
      Argumentss:
          new_state (torch.ByteTensor): The desired state
      This function is adapted from PyTorch repo (torch.cuda.set_rng_state)
      with a single change: the input state is not cloned. Cloning caused
      major performance issues for +4 GPU cases.
      """
      if hasattr(_C, "_cuda_setRNGState") and callable(
          _C._cuda_setRNGState
      ):
          # older PyTorch
          def cb():
              with device_ctx_manager(device):
                  _C._cuda_setRNGState(new_state)
  
      else:
          # newer PyTorch
          if device == -1:
              device = torch.device("cuda")
          elif isinstance(device, str):
              device = torch.device(device)
          elif isinstance(device, int):
              device = torch.device("cuda", device)
  
          def cb():
              idx = device.index
              if idx is None:
                  idx = torch.cuda.current_device()
              default_generator = torch.cuda.default_generators[idx]
              default_generator.set_state(new_state)
  
      _lazy_call(cb)
  
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 1 for calling a global function with example usage: _set_cuda_rng_state(new_state, device) without return types