code: def test_save_load(self):
      class Foo(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.foo = torch.nn.Linear(2, 2, bias=False)
              self.bar = torch.nn.Linear(2, 2, bias=False)
  
              self.reset_parameters()
  
          def reset_linear(self, module, generator):
              init.kaiming_uniform_(
                  module.weight, a=math.sqrt(5), generator=generator
              )
  
          def reset_parameters(self):
              generator = torch.Generator()
              generator.manual_seed(1)
              self.reset_linear(self.foo, generator)
  
              generator = torch.Generator()
              generator.manual_seed(2)
              self.reset_linear(self.bar, generator)
  
          def forward(self, x):
              x = self.foo(x)
              x = self.bar(x)
  
              generator = torch.Generator()
              generator.manual_seed(3)
              r = torch.empty_like(x)
              r.normal_(0.0, 1.0, generator=generator)
  
              return x, r
  
      eager_foo = Foo()
  
      script_module = torch.jit.script(Foo())
      saved_module = io.BytesIO()
      torch.jit.save(script_module, saved_module)
      saved_module.seek(0)
  
      loaded_module = torch.jit.load(saved_module)
  
      self.assertEqual(eager_foo.foo.weight, loaded_module.foo.weight)
      self.assertEqual(eager_foo.bar.weight, loaded_module.bar.weight)
  
      try:
          # Run this 3 times so make sure that the generator seed is being set
          # every time forward is called
          for i in range(3):
              x = torch.ones(2, 2)
              out1, r1 = eager_foo(x)
              out2, r2 = loaded_module(x)
  
              try:
                  self.assertEqual(out1, out2)
              except:  # noqa: B001, E722
                  print(f"Iteration {i}:\n{out1=}\n{out2=}")
                  raise
  
              try:
                  self.assertEqual(r1, r2)
              except:  # noqa: B001, E722
                  print(f"Iteration {i}:\n{r1=}\n{r2=}")
                  raise
      except:  # noqa: B001, E722
          print(loaded_module.forward.code)
          raise
  
file_path: pytorch/test/jit/test_generator.py
imports: import io
  import math
  import unittest
  import torch
  from torch.nn import init
  from torch.testing._internal.common_utils import skipIfLegacyJitExecutor
  from torch.testing._internal.jit_utils import JitTestCase
label_desc: Use image node_id 5 for calling the TestGenerator obj's underlying member method code with example usage: obj.test_save_load() and returns: x,  r