def __init__(
        self,
        input_dim: int = 768,
        hidden_dim: int = 2048,
        output_dim: int = 256,
    ):
        """Initializes the MSNProjectionHead with the specified dimensions."""
        super().__init__(
            blocks=[
                (input_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.GELU()),
                (hidden_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.GELU()),
                (hidden_dim, output_dim, None, None),
            ]
        )
Initializes the MSNProjectionHead with the specified dimensions.
lightly/models/modules/heads.py