code: def SwavPrototypesHeadFSDP(
      model_config: AttrDict,
      dims: List[int],
      use_bn: bool,
      num_clusters: int,
      use_bias: bool = True,
      return_embeddings: bool = True,
      skip_last_bn: bool = True,
      normalize_feats: bool = True,
  ):
      """
      SwAV head specific FSDP wrapping: we keep the full precision for the prototypes
  
      This is important for convergence:
      Since we "normalize" this layer in the update hook, we need to keep its
      weights in full precision. It is output is going into the loss and used
      for clustering, so we need to have that in full precision as well.
      """
  
      head = SwAVPrototypesHead(
          model_config=model_config,
          dims=dims,
          use_bn=use_bn,
          num_clusters=num_clusters,
          use_bias=use_bias,
          return_embeddings=return_embeddings,
          skip_last_bn=skip_last_bn,
          normalize_feats=normalize_feats,
      )
      head = fsdp_auto_wrap_bn(head)
  
      prototypes_fp32_fsdp_config = model_config.FSDP_CONFIG.copy()
      prototypes_fp32_fsdp_config["flatten_parameters"] = False
      prototypes_fp32_fsdp_config["mixed_precision"] = False
      prototypes_fp32_fsdp_config["fp32_reduce_scatter"] = False
      prototypes_fp32_fsdp_config["compute_dtype"] = torch.float32
      for j in range(head.nmb_heads):
          module = getattr(head, "prototypes" + str(j))
          module = fsdp_wrapper(module, **prototypes_fp32_fsdp_config)
          setattr(head, "prototypes" + str(j), module)
  
      return fsdp_wrapper(head, **model_config.FSDP_CONFIG)
  
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 1 for calling a global function with example usage: SwavPrototypesHeadFSDP(model_config, dims, use_bn, num_clusters, use_bias, return_embeddings, skip_last_bn, normalize_feats) and returns: fsdp_wrapper