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