code: def DINOHeadFSDP(
      model_config: AttrDict,
      in_dim: int,
      num_clusters: List[int],
      use_bn: bool = False,
      normalize_last_layer: bool = True,
      num_layers: int = 3,
      hidden_dim: int = 2048,
      bottleneck_dim: int = 256,
  ):
      head = DINOHead(
          model_config=model_config,
          in_dim=in_dim,
          num_clusters=num_clusters,
          use_bn=use_bn,
          normalize_last_layer=normalize_last_layer,
          num_layers=num_layers,
          hidden_dim=hidden_dim,
          bottleneck_dim=bottleneck_dim,
      )
  
      # TODO - FSDP does not work around "nn.utils.weight_norm" when
      #  required_grad is set to False on "weight_g", as it seems that
      #  FSDP requires all the parameters to be requiring grad or none
      #  of them, so we workaround it by making "weight_g" a buffer
      #  instead of a nn.Parameter
      if normalize_last_layer:
          head.prototypes0 = NormalizedLinearLayer(
              bottleneck_dim,
              num_clusters[0],
              init_weight_v=head.prototypes0.weight_v.data,
          )
  
      # Wrap prototypes in FP32
      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
      head.prototypes0 = fsdp_wrapper(
          head.prototypes0, **prototypes_fp32_fsdp_config
      )
  
      # Wrap the rest of the head
      return fsdp_wrapper(head, **model_config.FSDP_CONFIG)
  
file_path: vissl/vissl/models/heads/dino_head.py
imports: from typing import List, Optional
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.utils.weight_norm import WeightNorm
  from vissl.config import AttrDict
  from vissl.models.heads import register_model_head
  from vissl.models.model_helpers import trunc_normal_
  from vissl.utils.fsdp_utils import fsdp_wrapper
label_desc: Use image node_id 1 for calling a global function with example usage: DINOHeadFSDP(model_config, in_dim, num_clusters, use_bn, normalize_last_layer, num_layers, hidden_dim, bottleneck_dim) and returns: fsdp_wrapper