class_docstr: Helper to add to hooks to create and update a momentum encoder in the loss
class_docstr_tok: ["Helper","to","add","to","hooks","to","create","and","update","a","momentum","encoder","in","the","loss"]
docstr: EMA update: Each teacher parameter becomes a weighted average
  of its old self and the newest student.
docstr_tok: ["EMA","update",":","Each","teacher","parameter","becomes","a","weighted","average","of","its","old","self","and","the","newest","student","."]
code: def update_momentum_network(
      task: tasks.ClassyTask,
      init_teacher_momentum: float,
      with_cosine_schedule: bool = True,
  ) -> None:
      """
      EMA update: Each teacher parameter becomes a weighted average
      of its old self and the newest student.
      """
      with torch.no_grad():
          # Compute teacher momentum based on momentum schedule
          if with_cosine_schedule:
              m = 1 - 0.5 * (1 - init_teacher_momentum) * (
                  math.cos(
                      math.pi * task.iteration / task.max_iteration
                  )
                  + 1
              )
          else:
              m = init_teacher_momentum
  
          # Log the value of the momentum
          task.additional_log_data["teacher_momentum"] = m
  
          # Access the raw student and teacher models
          task_model = get_no_ddp_model(task.model)
          teacher_model = get_no_ddp_model(task.loss.momentum_teacher)
  
          # EMA update for the teacher parameters
          for param_q, param_k in zip(
              task_model.parameters(), teacher_model.parameters()
          ):
              param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)
  
file_path: vissl/vissl/hooks/hook_utils.py
imports: import logging
  import math
  import torch
  import torch.nn
  from classy_vision import tasks
  from torch.nn.parallel import DistributedDataParallel
  from vissl.config import AttrDict
  from vissl.models import build_model
  from vissl.models.model_helpers import get_no_ddp_model
  from vissl.utils.fsdp_utils import fsdp_recursive_reset_lazy_init, fsdp_wrapper
label_desc: Use image node_id 2 for calling the MomentumTeacherInLossHook obj's underlying member method code with example usage: obj.update_momentum_network(task, init_teacher_momentum, with_cosine_schedule) without return types