code: def _test_update_bn(self, dnn, dl_x, dl_xy, cuda):
      preactivation_sum = torch.zeros(dnn.n_features)
      preactivation_squared_sum = torch.zeros(dnn.n_features)
      if cuda:
          preactivation_sum = preactivation_sum.cuda()
          preactivation_squared_sum = preactivation_squared_sum.cuda()
      total_num = 0
      for x in dl_x:
          x = x[0]
          if cuda:
              x = x.cuda()
  
          dnn.forward(x)
          preactivations = dnn.compute_preactivation(x)
          if len(preactivations.shape) == 4:
              preactivations = preactivations.transpose(1, 3)
          preactivations = preactivations.contiguous().view(
              -1, dnn.n_features
          )
          total_num += preactivations.shape[0]
  
          preactivation_sum += torch.sum(preactivations, dim=0)
          preactivation_squared_sum += torch.sum(
              preactivations**2, dim=0
          )
  
      preactivation_mean = preactivation_sum / total_num
      preactivation_var = preactivation_squared_sum / total_num
      preactivation_var = preactivation_var - preactivation_mean**2
  
      update_bn(dl_xy, dnn, device=x.device)
      self.assertEqual(preactivation_mean, dnn.bn.running_mean)
      self.assertEqual(
          preactivation_var, dnn.bn.running_var, atol=1e-1, rtol=0
      )
  
      def _reset_bn(module):
          if issubclass(
              module.__class__, torch.nn.modules.batchnorm._BatchNorm
          ):
              module.running_mean = torch.zeros_like(
                  module.running_mean
              )
              module.running_var = torch.ones_like(module.running_var)
  
      # reset batch norm and run update_bn again
      dnn.apply(_reset_bn)
      update_bn(dl_xy, dnn, device=x.device)
      self.assertEqual(preactivation_mean, dnn.bn.running_mean)
      self.assertEqual(
          preactivation_var, dnn.bn.running_var, atol=1e-1, rtol=0
      )
      # using the dl_x loader instead of dl_xy
      dnn.apply(_reset_bn)
      update_bn(dl_x, dnn, device=x.device)
      self.assertEqual(preactivation_mean, dnn.bn.running_mean)
      self.assertEqual(
          preactivation_var, dnn.bn.running_var, atol=1e-1, rtol=0
      )
  
file_path: pytorch/test/optim/test_swa_utils.py
imports: import itertools
  import pickle
  import torch
  from torch.optim.swa_utils import AveragedModel, update_bn, get_swa_multi_avg_fn, get_ema_multi_avg_fn
  from torch.testing._internal.common_utils import TestCase, load_tests, parametrize, instantiate_parametrized_tests
label_desc: Use image node_id 8 for calling the TestSWAUtils obj's underlying member method code with example usage: obj._test_update_bn(dnn, dl_x, dl_xy, cuda) without return types