code: def test_new_spectral_norm_load_state_dict(self):
      for activate_times in (0, 3):
          inp = torch.randn(2, 3)
          m = nn.Linear(3, 5)
          snm = torch.nn.utils.parametrizations.spectral_norm(m)
          snm.train()
  
          for _ in range(activate_times):
              snm(inp)
  
          state_dict = deepcopy(snm.state_dict())
          self.assertEqual(
              {
                  "parametrizations.weight.original",
                  "bias",
                  "parametrizations.weight.0._v",
                  "parametrizations.weight.0._u",
              },
              set(state_dict.keys()),
          )
  
          # test that non-strict loading works
          non_strict_state_dict = deepcopy(state_dict)
          non_strict_state_dict["nonsense"] = "nonsense"
          with self.assertRaisesRegex(
              RuntimeError,
              r'Unexpected key\(s\) in state_dict: "nonsense"',
          ):
              snm.load_state_dict(non_strict_state_dict, strict=True)
          snm.load_state_dict(non_strict_state_dict, strict=False)
          del non_strict_state_dict["parametrizations.weight.original"]
          snm.load_state_dict(non_strict_state_dict, strict=False)
          del non_strict_state_dict["parametrizations.weight.0._u"]
          snm.load_state_dict(non_strict_state_dict, strict=False)
          del non_strict_state_dict["parametrizations.weight.0._v"]
          snm.load_state_dict(non_strict_state_dict, strict=False)
          non_strict_state_dict[
              "weight"
          ] = snm.weight.detach().clone()  # set W as a buffer
          snm.load_state_dict(non_strict_state_dict, strict=False)
          del non_strict_state_dict._metadata[
              "parametrizations.weight.0"
          ]  # remove metadata info
          snm.load_state_dict(non_strict_state_dict, strict=False)
          del non_strict_state_dict["weight"]  # remove W buffer
          snm.load_state_dict(non_strict_state_dict, strict=False)
          del non_strict_state_dict["bias"]
          snm.load_state_dict(non_strict_state_dict, strict=False)
  
          # normal state_dict
  
          # test that re-wrapping does not matter
          m = torch.nn.utils.parametrize.remove_parametrizations(
              snm, "weight"
          )
          snm = torch.nn.utils.parametrizations.spectral_norm(m)
  
          snm.load_state_dict(state_dict)
          with torch.no_grad():
              snm.eval()
              out0_eval = snm(inp)
              snm.train()
              out1_train = snm(inp)
              out2_train = snm(inp)
              snm.eval()
              out3_eval = snm(inp)
  
          # test that re-wrapping does not matter
          m = torch.nn.utils.parametrize.remove_parametrizations(
              snm, "weight"
          )
          snm = torch.nn.utils.parametrizations.spectral_norm(m)
  
          # Test normal loading
          snm.load_state_dict(state_dict)
          with torch.no_grad():
              snm.eval()
              self.assertEqual(out0_eval, snm(inp))
              snm.train()
              self.assertEqual(out1_train, snm(inp))
              self.assertEqual(out2_train, snm(inp))
              snm.eval()
              self.assertEqual(out3_eval, snm(inp))
  
imports: from copy import deepcopy
  from itertools import product
  import pickle
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.nn.init
  import torch.nn.utils.parametrize
  from torch.nn import Parameter
  from torch.testing._internal.common_utils import run_tests, skipIfNoLapack, TemporaryFileName, instantiate_parametrized_tests, set_default_dtype
  from torch.testing._internal.common_device_type import instantiate_device_type_tests
  from torch.testing._internal.common_cuda import TEST_MULTIGPU
  from torch.testing._internal.common_nn import NNTestCase
  from torch.testing._internal.common_utils import gradcheck
label_desc: Use image node_id 19 for calling the TestNNParametrization obj's underlying member method code with example usage: obj.test_new_spectral_norm_load_state_dict() without return types