code: def simclr_head(
      num_features: int = 2048,
      embedding_dim: int = 128,
      dims: List[int] = [2048],
      use_bn: bool = True,
      **kwargs,
  ) -> nn.Module:
      cfg = VISSLAdapter.get_model_config_template()
      head_kwargs = {
          "dims": [num_features] + dims + [embedding_dim],
          "use_bn": use_bn,
      }
  
      cfg.HEAD.PARAMS.append(["simclr_head", head_kwargs])
  
      head = MODEL_HEADS_REGISTRY["simclr_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 1 for calling a global function with example usage: simclr_head(num_features, embedding_dim, dims, use_bn) and returns: head