code: def test_averaged_model_state_dict(self) -> None:
      dnn = torch.nn.Sequential(
          torch.nn.Conv2d(1, 5, kernel_size=3), torch.nn.Linear(5, 10)
      )
      averaged_dnn = AveragedModel(dnn)
      averaged_dnn2 = AveragedModel(dnn)
      n_updates = 10
      for _ in range(n_updates):
          for p in dnn.parameters():
              p.detach().add_(torch.randn_like(p))
          averaged_dnn.update_parameters(dnn)
      averaged_dnn2.load_state_dict(averaged_dnn.state_dict())
      for p_swa, p_swa2 in zip(
          averaged_dnn.parameters(), averaged_dnn2.parameters()
      ):
          torch.testing.assert_close(p_swa, p_swa2, check_device=False)
      self.assertTrue(
          averaged_dnn.n_averaged == averaged_dnn2.n_averaged
      )
  
file_path: tnt/tests/utils/test_swa.py
imports: import itertools
  import unittest
  from copy import deepcopy
  from typing import List, Tuple
  import torch
  from torchtnt.utils.swa import AveragedModel
label_desc: Use image node_id 4 for calling the TestSWA obj's underlying member method code with example usage: obj.test_averaged_model_state_dict() without return types