class_docstr: This hook corresponds to the loss proposed in the "Momentum Contrast
  for Unsupervised Visual Representation Learning" paper, from Kaiming He et al.
  See http://arxiv.org/abs/1911.05722 for details
  and https://github.com/facebookresearch/moco for a reference implementation,
  reused here.
  
  Called after every forward pass to update the momentum encoder. At the beginning
  of training i.e. after 1st forward call, the encoder is contructed and updated.
docstr: - Update the momentum encoder.
  - Compute the key reusing the updated moco-encoder. If we use the
    batch shuffling, the perform global shuffling of the batch
    and then run the moco encoder to compute the features.
    We unshuffle the computer features and use the features
    as "key" in computing the moco loss.
docstr_tok: ["-","Update","the","momentum","encoder",".","-","Compute","the","key","reusing","the","updated","moco-encoder",".","If","we","use","the","batch","shuffling",",","the","perform","global","shuffling","of","the","batch","and","then","run","the","moco","encoder","to","compute","the","features",".","We","unshuffle","the","computer","features","and","use","the","features","as","``","key","''","in","computing","the","moco","loss","."]
code: def on_forward(self, task: tasks.ClassyTask) -> None:
      """
      - Update the momentum encoder.
      - Compute the key reusing the updated moco-encoder. If we use the
        batch shuffling, the perform global shuffling of the batch
        and then run the moco encoder to compute the features.
        We unshuffle the computer features and use the features
        as "key" in computing the moco loss.
      """
  
      # Update the momentum encoder
      if task.loss.moco_encoder is None:
          self._build_moco_encoder(task)
          self.is_distributed = is_distributed_training_run()
          logging.info("MoCo: Distributed setup, shuffling batches")
      else:
          self._update_momentum_encoder(task)
  
      # Compute key features. We do not backpropagate in this codepath
      im_k = task.last_batch.sample["data_momentum"][0]
  
      if self.is_distributed and self.shuffle_batch:
          # shuffle for making use of BN
          im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k, task)
  
      k = task.loss.moco_encoder(im_k)[0]
      k = torch.nn.functional.normalize(k, dim=1)
  
      if self.is_distributed and self.shuffle_batch:
          # undo shuffle
          k = self._batch_unshuffle_ddp(k, idx_unshuffle)
  
      task.loss.key = k
  
file_path: vissl/vissl/hooks/moco_hooks.py
imports: import logging
  import torch
  from classy_vision import tasks
  from classy_vision.generic.distributed_util import is_distributed_training_run
  from classy_vision.hooks.classy_hook import ClassyHook
  from vissl.models import build_model
  from vissl.utils.env import get_machine_local_and_dist_rank
  from vissl.utils.misc import concat_all_gather
label_desc: Use image node_id 6 for calling the MoCoHook obj's underlying member method code with example usage: obj.on_forward(task) without return types