code: def swav_head(
      num_features: int = 2048,
      embedding_dim: int = 128,
      dims: List[int] = [2048],
      use_bn: bool = True,
      num_clusters: Union[int, List[int]] = [3000],
      use_bias: bool = True,
      return_embeddings: bool = True,
      skip_last_bn: bool = True,
      normalize_feats: bool = True,
      activation_name: str = "ReLU",
      use_weight_norm_prototypes: bool = False,
      **kwargs,
  ) -> nn.Module:
      cfg = VISSLAdapter.get_model_config_template()
      head_kwargs = {
          "dims": [num_features] + dims + [embedding_dim],
          "use_bn": use_bn,
          "num_clusters": [num_clusters]
          if isinstance(num_clusters, int)
          else num_clusters,
          "use_bias": use_bias,
          "return_embeddings": return_embeddings,
          "skip_last_bn": skip_last_bn,
          "normalize_feats": normalize_feats,
          "activation_name": activation_name,
          "use_weight_norm_prototypes": use_weight_norm_prototypes,
      }
  
      cfg.HEAD.PARAMS.append(["swav_head", head_kwargs])
  
      head = MODEL_HEADS_REGISTRY["swav_head"](cfg, **head_kwargs)
      head.model_config = cfg
  
      return head
  
file_path: lightning-flash/src/flash/image/embedding/heads/vissl_heads.py
imports: from typing import List, Union
  import torch.nn
  from torch import Tensor
  from flash.core.registry import FlashRegistry
  from flash.core.utilities.imports import _VISSL_AVAILABLE
label_desc: Use image node_id 2 for calling a global function with example usage: swav_head(num_features, embedding_dim, dims, use_bn, num_clusters, use_bias, return_embeddings, skip_last_bn, normalize_feats, activation_name, use_weight_norm_prototypes) and returns: head