class_docstr: L2 Normalize the prototypes in swav training. Optional.
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_loss":
          return
      if not task.config.LOSS["swav_loss"].normalize_last_layer:
          return
      with torch.no_grad():
          try:
              # This is either single GPU model or a FSDP.
              assert len(task.model.heads) == 1
              for j in range(task.model.heads[0].nmb_heads):
                  module = getattr(
                      task.model.heads[0], "prototypes" + str(j)
                  )
                  # Determine the context we need to use. For FSDP, we
                  # need the summon_full_params context, which ensures that
                  # full weights for this layer is all_gathered and after
                  # normalization, the changes are persisted in the local
                  # shards. All ranks do the same normalization, so all
                  # changes should be saved.
                  ctx = contextlib.suppress()
                  if hasattr(module, "summon_full_params"):
                      ctx = module.summon_full_params()
                  with ctx:
                      w = module.weight.data.clone()
                      w = nn.functional.normalize(w, dim=1, p=2)
                      module.weight.copy_(w)
          except AttributeError:
              # This is a DDP wrapped one.
              assert len(task.model.module.heads) == 1
              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)
  
file_path: vissl/vissl/hooks/swav_hooks.py
imports: import contextlib
  import torch
  import torch.nn
  from classy_vision import tasks
  from classy_vision.hooks.classy_hook import ClassyHook
label_desc: Use image node_id 1 for calling the NormalizePrototypesHook obj's underlying member method code with example usage: obj.on_update(task) without return types