class_docstr: Hook which keeps a track of the exponential moving average (EMA) of the model's
  parameters and applies the EMA params to the model during the test phases.
  
  Saving the state in cpu will save gpu memory, but will make training slower since
  the model parameters will need to be moved to cpu before the averaging.
  
  Note:
      This hooks stores two additional copies of the model's parameters, which will
      increase memory usage significantly.
class_docstr_tok: ["Hook","which","keeps","a","track","of","the","exponential","moving","average","(","EMA",")","of","the","model's","parameters","and","applies","the","EMA","params","to","the","model","during","the","test","phases",".","Saving","the","state","in","cpu","will","save","gpu","memory",",","but","will","make","training","slower","since","the","model","parameters","will","need","to","be","moved","to","cpu","before","the","averaging",".","Note",":","This","hooks","stores","two","additional","copies","of","the","model","'s","parameters",",","which","will","increase","memory","usage","significantly","."]
code: def on_step(self, task) -> None:
      if not task.train:
          return
  
      with torch.no_grad():
          if self.use_optimization(task):
              torch._foreach_mul_(self.ema_model_state_list, self.decay)
              torch._foreach_add_(
                  self.ema_model_state_list,
                  self.param_list,
                  alpha=(1 - self.decay),
              )
          else:
              for name, param in self.get_model_state_iterator(
                  task.base_model
              ):
                  self.state.ema_model_state[
                      name
                  ] = self.decay * self.state.ema_model_state[name] + (
                      1 - self.decay
                  ) * param.to(
                      device=self.device
                  )
  
file_path: ClassyVision/classy_vision/hooks/exponential_moving_average_model_hook.py
imports: import itertools
  import logging
  from typing import Any, Dict, Iterable, Tuple
  import torch
  import torch.nn
  from classy_vision.generic.util import get_torch_version
  from classy_vision.hooks import register_hook
  from classy_vision.hooks.classy_hook import ClassyHook
label_desc: Use image node_id 7 for calling the ExponentialMovingAverageModelHook obj's underlying member method code with example usage: obj.on_step(task) without return types