code: def _test_averaged_model_exponential(
      self, use_buffers: bool, skip_deepcopy: bool
  ) -> None:
      # Test AveragedModel with EMA as avg_fn and use_buffers as True.
      dnn = torch.nn.Sequential(
          torch.nn.Conv2d(1, 5, kernel_size=3),
          torch.nn.BatchNorm2d(5, momentum=0.3),
          torch.nn.Linear(5, 10),
      )
      decay: float = 0.9
  
      averaged_dnn = AveragedModel(
          deepcopy(dnn) if skip_deepcopy else dnn,
          averaging_method="ema",
          ema_decay=decay,
          use_buffers=use_buffers,
          skip_deepcopy=skip_deepcopy,
      )
  
      if use_buffers:
          dnn_params = list(
              itertools.chain(dnn.parameters(), dnn.buffers())
          )
      else:
          dnn_params = list(dnn.parameters())
  
      averaged_params = [
          torch.zeros_like(param)
          for param in dnn_params
          if param.size() != torch.Size([])
      ]
  
      n_updates = 10
      for i in range(n_updates):
          updated_averaged_params = []
          for p, p_avg in zip(dnn_params, averaged_params):
              if p.size() == torch.Size([]):
                  continue
              p.detach().add_(torch.randn_like(p))
              if i == 0:
                  updated_averaged_params.append(p.clone())
              else:
                  updated_averaged_params.append(
                      (p_avg * decay + p * (1 - decay)).clone()
                  )
          averaged_dnn.update_parameters(dnn)
          averaged_params = updated_averaged_params
  
      if use_buffers:
          for p_avg, p_swa in zip(
              averaged_params,
              itertools.chain(
                  averaged_dnn.module.parameters(),
                  averaged_dnn.module.buffers(),
              ),
          ):
              torch.testing.assert_close(
                  p_avg, p_swa, check_device=False
              )
      else:
          for p_avg, p_swa in zip(
              averaged_params, averaged_dnn.parameters()
          ):
              torch.testing.assert_close(
                  p_avg, p_swa, check_device=False
              )
          for b_avg, b_swa in zip(
              dnn.buffers(), averaged_dnn.module.buffers()
          ):
              torch.testing.assert_close(
                  b_avg, b_swa, check_device=False
              )
  
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 6 for calling the TestSWA obj's underlying member method code with example usage: obj._test_averaged_model_exponential(use_buffers, skip_deepcopy) without return types