code: def _test_exponential_moving_average_hook(
      self, model_device, hook_device
  ):
      task = mock.MagicMock()
      model = TestModel().to(
          device=self._map_device_string(model_device)
      )
      task.base_model = model
      task.train = True
      decay = 0.5
      num_updates = 10
      model.init_fc_weight()
      exponential_moving_average_hook = (
          ExponentialMovingAverageModelHook(
              decay=decay, device=hook_device
          )
      )
  
      exponential_moving_average_hook.on_start(task)
      exponential_moving_average_hook.on_phase_start(task)
      # set the weights to all ones and simulate 10 updates
      task.base_model.update_fc_weight()
      fc_weight = model.fc.weight.clone()
      for _ in range(num_updates):
          exponential_moving_average_hook.on_step(task)
      exponential_moving_average_hook.on_phase_end(task)
      # the model weights shouldn't have changed
      self.assertTrue(torch.allclose(model.fc.weight, fc_weight))
  
      # simulate a test phase now
      task.train = False
      exponential_moving_average_hook.on_phase_start(task)
      exponential_moving_average_hook.on_phase_end(task)
  
      # the model weights should be updated to the ema weights
      self.assertTrue(
          torch.allclose(
              model.fc.weight,
              fc_weight * (1 - math.pow(1 - decay, num_updates)),
          )
      )
  
      # simulate a train phase again
      task.train = True
      exponential_moving_average_hook.on_phase_start(task)
  
      # the model weights should be back to the old value
      self.assertTrue(torch.allclose(model.fc.weight, fc_weight))
  
imports: import copy
  import math
  import unittest
  import unittest.mock
  import torch
  import torch.nn
  from classy_vision.hooks import ExponentialMovingAverageModelHook
  from classy_vision.models import ClassyModel
  from test.generic.hook_test_utils import HookTestBase
label_desc: Use image node_id 2 for calling the TestExponentialMovingAverageModelHook obj's underlying member method code with example usage: obj._test_exponential_moving_average_hook(model_device, hook_device) without return types