class_docstr: This hook is for the extension of the SwAV loss proposed in paper
  https://arxiv.org/abs/2006.09882 by Caron et al. The loss combines
  the benefits of using the SwAV approach with the momentum encoder
  as used in MoCo.
class_docstr_tok: ["This","hook","is","for","the","extension","of","the","SwAV","loss","proposed","in","paper","https",":","\/\/arxiv.org\/abs\/2006.09882","by","Caron","et","al",".","The","loss","combines","the","benefits","of","using","the","SwAV","approach","with","the","momentum","encoder","as","used","in","MoCo","."]
docstr: Forward pass with momentum network. We forward momentum encoder
  only on the single resolution crops that are used for assignment
  in the swav loss.
docstr_tok: ["Forward","pass","with","momentum","network",".","We","forward","momentum","encoder","only","on","the","single","resolution","crops","that","are","used","for","assignment","in","the","swav","loss","."]
code: def on_forward(self, task: tasks.ClassyTask) -> None:
      """
      Forward pass with momentum network. We forward momentum encoder
      only on the single resolution crops that are used for assignment
      in the swav loss.
      """
  
      # Update the momentum encoder
      if task.loss.momentum_encoder is None:
          self._build_momentum_network(task)
      else:
          self._update_momentum_network(task)
  
      if task.loss.num_iteration >= self.momentum_eval_mode_iter_start:
          task.loss.momentum_encoder.eval()
          if (
              task.loss.num_iteration
              == self.momentum_eval_mode_iter_start
          ):
              logging.info(
                  "Momentum network will be used in eval mode."
              )
      else:
          task.loss.momentum_encoder.train()
  
      # Compute momentum features. We do not backpropagate in this codepath
      im_k = [
          task.last_batch.sample["input"][i]
          for i in self.crops_for_assign
      ]
      output = task.loss.momentum_encoder(im_k)[0]
      task.loss.momentum_scores = output[1:]
      task.loss.momentum_embeddings = output[0]
  
file_path: vissl/vissl/hooks/swav_momentum_hooks.py
imports: import logging
  from typing import List
  import torch
  from classy_vision import tasks
  from classy_vision.generic.distributed_util import init_distributed_data_parallel_model
  from classy_vision.hooks.classy_hook import ClassyHook
  from torch import nn
  from vissl.models import build_model
  from vissl.utils.env import get_machine_local_and_dist_rank
label_desc: Use image node_id 4 for calling the SwAVMomentumHook obj's underlying member method code with example usage: obj.on_forward(task) without return types