class_docstr: L2 Normalize the prototypes in swav training. Optional.
  We normalize the momentum_encoder output prototypes as well
  additionally.
class_docstr_tok: ["L2","Normalize","the","prototypes","in","swav","training",".","Optional",".","We","normalize","the","momentum_encoder","output","prototypes","as","well","additionally","."]
docstr: Optionally normalize prototypes
docstr_tok: ["Optionally","normalize","prototypes"]
code: def on_update(self, task: "tasks.ClassyTask") -> None:
      """
      Optionally normalize prototypes
      """
      if not task.config["LOSS"]["name"] == "swav_momentum_loss":
          return
      if not task.config.LOSS[
          "swav_momentum_loss"
      ].normalize_last_layer:
          return
      with torch.no_grad():
          try:
              for j in range(task.model.heads[0].nmb_heads):
                  w = getattr(
                      task.model.heads[0], "prototypes" + str(j)
                  ).weight.data.clone()
                  w = nn.functional.normalize(w, dim=1, p=2)
                  getattr(
                      task.model.heads[0], "prototypes" + str(j)
                  ).weight.copy_(w)
          except AttributeError:
              # TODO (mathildecaron): don't use getattr
              for j in range(task.model.module.heads[0].nmb_heads):
                  w = getattr(
                      task.model.module.heads[0], "prototypes" + str(j)
                  ).weight.data.clone()
                  w = nn.functional.normalize(w, dim=1, p=2)
                  getattr(
                      task.model.module.heads[0], "prototypes" + str(j)
                  ).weight.copy_(w)
          try:
              for j in range(
                  task.loss.momentum_encoder.heads[0].nmb_heads
              ):
                  w = getattr(
                      task.loss.momentum_encoder.heads[0],
                      "prototypes" + str(j),
                  ).weight.data.clone()
                  w = nn.functional.normalize(w, dim=1, p=2)
                  getattr(
                      task.loss.momentum_encoder.heads[0],
                      "prototypes" + str(j),
                  ).weight.copy_(w)
          except AttributeError:
              for j in range(
                  task.loss.momentum_encoder.module.heads[0].nmb_heads
              ):
                  w = getattr(
                      task.loss.momentum_encoder.module.heads[0],
                      "prototypes" + str(j),
                  ).weight.data.clone()
                  w = nn.functional.normalize(w, dim=1, p=2)
                  getattr(
                      task.loss.momentum_encoder.module.heads[0],
                      "prototypes" + str(j),
                  ).weight.copy_(w)
  
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 1 for calling the SwAVMomentumNormalizePrototypesHook obj's underlying member method code with example usage: obj.on_update(task) without return types