code: def test_bn_stats(self):
      base_self = self
  
      class TestHook(ClassyHook):
          on_start = ClassyHook._noop
          on_phase_start = ClassyHook._noop
          on_phase_end = ClassyHook._noop
          on_end = ClassyHook._noop
  
          def __init__(self):
              self.train_bn_stats = None
              self.test_bn_stats = None
  
          def on_step(self, task):
              if task.train:
                  self.train_bn_stats = base_self._get_bn_stats(
                      task.base_model
                  )
              else:
                  self.test_bn_stats = base_self._get_bn_stats(
                      task.base_model
                  )
  
      config = get_test_mlp_task_config()
      task = build_task(config)
      num_samples = 10
      precise_batch_norm_hook = PreciseBatchNormHook(num_samples)
      test_hook = TestHook()
      task.set_hooks([precise_batch_norm_hook, test_hook])
      trainer = ClassyTrainer()
      trainer.train(task)
  
      updated_bn_stats = self._get_bn_stats(task.base_model)
  
      # the stats should be modified after train steps but not after test steps
      self.assertFalse(
          self._compare_bn_stats(
              test_hook.train_bn_stats, updated_bn_stats
          )
      )
      self.assertTrue(
          self._compare_bn_stats(
              test_hook.test_bn_stats, updated_bn_stats
          )
      )
  
file_path: ClassyVision/test/hooks_precise_batch_norm_hook_test.py
imports: import copy
  import torch
  import torch.nn
  from classy_vision.hooks import ClassyHook
  from classy_vision.hooks.precise_batch_norm_hook import PreciseBatchNormHook
  from classy_vision.tasks import build_task
  from classy_vision.trainer import ClassyTrainer
  from test.generic.config_utils import get_test_mlp_task_config
  from test.generic.hook_test_utils import HookTestBase
label_desc: Use image node_id 5 for calling the TestPreciseBatchNormHook obj's underlying member method code with example usage: obj.test_bn_stats() without return types