code: def _run_averaged_steps(
      self, dnn: torch.nn.Module, swa_device: torch.device, ema: bool
  ) -> Tuple[List[torch.Tensor], torch.nn.Module]:
      ema_decay = 0.999
      if ema:
          averaged_dnn = AveragedModel(
              dnn,
              device=swa_device,
              averaging_method="ema",
              ema_decay=ema_decay,
          )
      else:
          averaged_dnn = AveragedModel(
              dnn, device=swa_device, averaging_method="swa"
          )
  
      averaged_params = [
          torch.zeros_like(param) for param in dnn.parameters()
      ]
  
      n_updates = 10
      for i in range(n_updates):
          for p, p_avg in zip(dnn.parameters(), averaged_params):
              p.detach().add_(torch.randn_like(p))
              if ema:
                  p_avg += (
                      p.detach()
                      * ema_decay ** (n_updates - i - 1)
                      * ((1 - ema_decay) if i > 0 else 1.0)
                  )
              else:
                  p_avg += p.detach() / n_updates
          averaged_dnn.update_parameters(dnn)
  
      return averaged_params, averaged_dnn
  
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 2 for calling the TestSWA obj's underlying member method code with example usage: obj._run_averaged_steps(dnn, swa_device, ema) and returns: averaged_params,  averaged_dnn