code: def test_bn_update(self):
      def test(net_cls, x_shape, y_shape, device):
          x = torch.rand(x_shape, device=device)
          y = torch.rand(y_shape, device=device)
  
          dnn = net_cls().to(device)
          orig_momentum = dnn.bn.momentum
          dnn.train()
          self._test_bn_update(x, dnn, device)
          self._test_bn_update(x, dnn, device, label_tensor=y)
          self.assertTrue(dnn.training)
  
          # check that bn_update preserves eval mode
          dnn.eval()
          self._test_bn_update(x, dnn, device)
          self.assertFalse(dnn.training)
  
          # check that momentum is preserved
          assert equal(dnn.bn.momentum, orig_momentum)
  
      # Test bn_update for fully-connected and convolutional networks with
      # BatchNorm1d and BatchNorm2d respectively
      objects = 100
      input_features = 5
  
      class DNN(nn.Module):
          def __init__(self):
              super(DNN, self).__init__()
              self.n_features = 100
              self.fc1 = nn.Linear(input_features, self.n_features)
              self.bn = nn.BatchNorm1d(self.n_features)
  
          def compute_preactivation(self, x):
              return self.fc1(x)
  
          def forward(self, x):
              x = self.fc1(x)
              x = self.bn(x)
              return x
  
      test(DNN, (objects, input_features), objects, "cpu")
      if torch.cuda.is_available():
          test(DNN, (objects, input_features), objects, "cuda")
  
      # Test bn_update for convolutional network and BatchNorm2d
      objects = 100
      channels = 3
      height, width = 5, 5
  
      class CNN(nn.Module):
          def __init__(self):
              super(CNN, self).__init__()
              self.n_features = 10
              self.conv1 = nn.Conv2d(
                  channels, self.n_features, kernel_size=3, padding=1
              )
              self.bn = nn.BatchNorm2d(self.n_features, momentum=0.3)
  
          def compute_preactivation(self, x):
              return self.conv1(x)
  
          def forward(self, x):
              x = self.conv1(x)
              x = self.bn(x)
              return x
  
      test(CNN, (objects, channels, height, width), objects, "cpu")
      if torch.cuda.is_available():
          test(CNN, (objects, channels, height, width), objects, "cuda")
  
imports: import contextlib
  import functools
  import re
  import warnings
  from unittest import TestCase
  import torch
  from pytext.optimizer import Lamb, StochasticWeightAveraging
  from torch import nn, optim, sparse
  from torch.autograd import Variable
  from torch.utils import data
label_desc: Use image node_id 18 for calling the TestSWA obj's underlying member method code with example usage: obj.test_bn_update() and returns: self, x, self, x