def __init__(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dim: int = 2048,
        num_blocks: int = 12,
        mlp_ratio: float = 4.0,
        proj_drop: float = 0.0,
        act_layer: Type[Module] = GELU,
        norm_layer: Type[Module] = LayerNorm,
        mlp_layer: Type[Module] = Mlp,
        block_fn: Type[Module] = AIMPredictionHeadBlock,
    ) -> None:
        """Initializes the AIMPredictionHead module with the specified parameters."""

        super().__init__()
        self.blocks = Sequential(
            # Linear layer to project the input dimension to the hidden dimension.
            Linear(input_dim, hidden_dim, bias=False),
            # Main blocks.
            *[
                block_fn(  # type: ignore[call-arg]
                    input_dim=hidden_dim,
                    output_dim=hidden_dim,
                    mlp_ratio=mlp_ratio,
                    proj_drop=proj_drop,
                    norm_layer=norm_layer,
                    act_layer=act_layer,
                    mlp_layer=mlp_layer,
                )
                for _ in range(num_blocks)
            ],
            # Linear layer to project the hidden dimension to the output dimension.
            norm_layer(hidden_dim),  # type: ignore[call-arg]
            Linear(hidden_dim, output_dim, bias=False),
        )
Initializes the AIMPredictionHead module with the specified parameters.
lightly/models/modules/heads_timm.py