def METHOD_NAME(
    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