code: def test_assign_batchnorm():
      x = np.random.standard_normal((10, 64))
      y = np.random.standard_normal((10, 1))
  
      model = tf.keras.Sequential()
      model.add(tf.keras.layers.Dense(16, activation="relu"))
      model.add(tf.keras.layers.BatchNormalization())
      model.add(tf.keras.layers.Dense(1))
  
      if hasattr(tf.keras.optimizers, "legacy"):
          opt = SWA(tf.keras.optimizers.legacy.SGD())
      else:
          opt = SWA(tf.keras.optimizers.SGD())
      model.compile(optimizer=opt, loss="mean_squared_error")
      model.fit(x, y, epochs=1)
  
      opt.assign_average_vars(model.variables)
      fit_bn(model, x, y)
  
file_path: addons/tensorflow_addons/optimizers/tests/stochastic_weight_averaging_test.py
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.optimizers import stochastic_weight_averaging
  from tensorflow_addons.optimizers.utils import fit_bn
label_desc: Use image node_id 5 for calling a global function with example usage: test_assign_batchnorm() without return types