def __init__(
        self,
        input_dim: int = 2048,
        hidden_dim: int = 8192,
        output_dim: int = 8192,
        num_layers: int = 3,
    ):
        """Initializes the VICRegProjectionHead with the specified dimensions.

        Args:
            input_dim:
                Dimensionality of the input features.
            hidden_dim:
                Dimensionality of the hidden layers.
            output_dim:
                Dimensionality of the output features.
            num_layers:
                Number of layers in the projection head.
        """
        hidden_layers = [
            (hidden_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.ReLU())
            for _ in range(num_layers - 2)  # Exclude first and last layer.
        ]
        super(VICRegProjectionHead, self).__init__(
            [
                (input_dim, hidden_dim, nn.BatchNorm1d(hidden_dim), nn.ReLU()),
                *hidden_layers,
                (hidden_dim, output_dim, None, None),
            ]
        )
Initializes the VICRegProjectionHead with the specified dimensions.

        Args:
            input_dim:
                Dimensionality of the input features.
            hidden_dim:
                Dimensionality of the hidden layers.
            output_dim:
                Dimensionality of the output features.
            num_layers:
                Number of layers in the projection head.
        
lightly/models/modules/heads.py