def __init__(
        self,
        input_dim: int,
        output_dim: int,
        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,
    ) -> None:
        """Initializes the AIMPredictionHeadBlock module with the specified parameters."""

        super().__init__()
        self.norm = norm_layer(input_dim)  # type: ignore[call-arg]
        self.mlp = mlp_layer(  # type: ignore[call-arg]
            in_features=input_dim,
            hidden_features=int(input_dim * mlp_ratio),
            out_features=output_dim,
            act_layer=act_layer,
            drop=proj_drop,
            bias=False,
        )
Initializes the AIMPredictionHeadBlock module with the specified parameters.
lightly/models/modules/heads_timm.py