code: def freeze_rng_state():
      rng_state = torch.get_rng_state()
      if torch.cuda.is_available():
          cuda_rng_state = torch.cuda.get_rng_state()
      yield
      if torch.cuda.is_available():
          torch.cuda.set_rng_state(cuda_rng_state)
      torch.set_rng_state(rng_state)
  
imports: import contextlib
  import numbers
  import os
  import re
  import warnings
  from itertools import product
  import numpy
  import pytest
  import torch
  import torch.cuda
  from numpy.testing import assert_allclose
  from pytest import approx
label_desc: Use image node_id 8 for calling a global function with example usage: freeze_rng_state() without return types