class_docstr: SwAV head used in https://arxiv.org/pdf/2006.09882.pdf paper.
  
  The head is composed of 2 parts
      1) projection of features to lower dimension like 128
      2) feature classification into clusters (also called prototypes)
  
  The projected features are L2 normalized before clustering step.
  
  Input: 2D torch.tensor of shape (N x C)
  
  Output: List(2D torch.tensor of shape N x num_clusters)
docstr: Args:
      batch (4D torch.tensor): shape (N x C x H x W)
  Returns:
      List(2D torch.tensor of shape N x num_clusters)
docstr_tok: ["Args",":","batch","(","4D","torch.tensor",")",":","shape","(","N","x","C","x","H","x","W",")","Returns",":","List","(","2D","torch.tensor","of","shape","N","x","num_clusters",")"]
code: def forward(self, batch: torch.Tensor):
      """
      Args:
          batch (4D torch.tensor): shape (N x C x H x W)
      Returns:
          List(2D torch.tensor of shape N x num_clusters)
      """
      batch = self.projection_head(batch)
  
      if self.normalize_feats:
          batch = nn.functional.normalize(batch, dim=1, p=2)
  
      out = []
      if self.return_embeddings:
          out.append(batch)
      if self.nmb_heads > 0:
          for i in range(self.nmb_heads):
              out.append(getattr(self, "prototypes" + str(i))(batch))
      return out
  
file_path: vissl/vissl/models/heads/swav_prototypes_head.py
imports: from typing import List
  import torch
  import torch.nn
  from vissl.config import AttrDict
  from vissl.models.heads import register_model_head
  from vissl.utils.fsdp_utils import fsdp_auto_wrap_bn, fsdp_wrapper
label_desc: Use image node_id 2 for calling the SwAVPrototypesHead obj's underlying member method code with example usage: obj.forward(batch) and returns: out