code: def _test_averaged_model(
      self,
      net_device: torch.device,
      swa_device: torch.device,
      ema: bool,
  ) -> None:
      dnn = torch.nn.Sequential(
          torch.nn.Conv2d(1, 5, kernel_size=3),
          torch.nn.ReLU(),
          torch.nn.MaxPool2d(kernel_size=2),
          torch.nn.BatchNorm2d(5, momentum=0.3),
          torch.nn.Conv2d(5, 2, kernel_size=3),
          torch.nn.ReLU(),
          torch.nn.Linear(5, 5),
          torch.nn.ReLU(),
          torch.nn.Linear(5, 10),
      ).to(net_device)
  
      averaged_params, averaged_dnn = self._run_averaged_steps(
          dnn, swa_device, ema
      )
  
      for p_avg, p_swa in zip(
          averaged_params, averaged_dnn.parameters()
      ):
          torch.testing.assert_close(p_avg, p_swa, check_device=False)
          # Check that AveragedModel is on the correct device
          self.assertTrue(p_swa.device == swa_device)
          self.assertTrue(p_avg.device == net_device)
      self.assertTrue(averaged_dnn.n_averaged.device == swa_device)
  
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 1 for calling the TestSWA obj's underlying member method code with example usage: obj._test_averaged_model(net_device, swa_device, ema) without return types