class_docstr: Hook to compute precise batch normalization statistics.
  
  Batch norm stats are calculated and updated during training, when the weights are
  also changing, which makes the calculations imprecise. This hook recomputes the
  batch norm stats at the end of a train phase to make them more precise. See
  `fvcore's documentation <https://github.com/facebookresearch/fvcore/blob/master/
  fvcore/nn/precise_bn.py>`_ for more information.
class_docstr_tok: ["Hook","to","compute","precise","batch","normalization","statistics",".","Batch","norm","stats","are","calculated","and","updated","during","training",",","when","the","weights","are","also","changing",",","which","makes","the","calculations","imprecise",".","This","hook","recomputes","the","batch","norm","stats","at","the","end","of","a","train","phase","to","make","them","more","precise",".","See","`","fvcore","'s","documentation","<","https",":","\/\/github.com\/facebookresearch\/fvcore\/blob\/master\/","fvcore\/nn\/precise_bn.py",">","`","_","for","more","information","."]
code: def on_step(self, task) -> None:
      if not task.train:
          return
  
      if self.cache_samples:
          if self.current_samples >= self.num_samples:
              return
          sample = recursive_copy_to_device(
              task.last_batch.sample,
              non_blocking=True,
              device=torch.device("cpu"),
          )
          self.cache.append(sample)
          self.current_samples += get_batchsize_per_replica(sample)
  
      else:
          if self.batch_size is not None:
              return
  
          self.batch_size = get_batchsize_per_replica(
              task.last_batch.sample["input"]
          )
  
file_path: ClassyVision/classy_vision/hooks/precise_batch_norm_hook.py
imports: import logging
  import math
  import torch
  from classy_vision.generic.util import get_batchsize_per_replica, recursive_copy_to_device, recursive_copy_to_gpu
  from classy_vision.hooks import ClassyHook, register_hook
  from fvcore.nn.precise_bn import update_bn_stats
label_desc: Use image node_id 5 for calling the PreciseBatchNormHook obj's underlying member method code with example usage: obj.on_step(task) without return types