def __init__(
        self,
        input_dim: int = 2048,
        hidden_dim: int = 8192,
        output_dim: int = 512,
        num_layers: int = 2,
        batch_norm: bool = True,
        use_bias: bool = False,
    ):
        """Initialize a new MMCRProjectionHead instance.

        Args:
            input_dim:
                Number of input dimensions.
            hidden_dim:
                Number of hidden dimensions.
            output_dim:
                Number of output dimensions.
            num_layers:
                Number of hidden layers.
            batch_norm:
                Whether or not to use batch norms.
            use_bias:
                Whether or not to use bias in the linear layers.
        """
        layers: List[
            Tuple[int, int, Optional[nn.Module], Optional[nn.Module], bool]
        ] = []

        # Add the first layer
        layers.append(
            (
                input_dim,
                hidden_dim,
                nn.BatchNorm1d(hidden_dim) if batch_norm else None,
                nn.ReLU(),
                use_bias,
            )
        )

        # Add the hidden layers
        for _ in range(num_layers - 1):
            layers.append(
                (
                    hidden_dim,
                    hidden_dim,
                    nn.BatchNorm1d(hidden_dim) if batch_norm else None,
                    nn.ReLU(),
                    use_bias,
                )
            )

        # Add the output layer
        layers.append((hidden_dim, output_dim, None, None, use_bias))
        super().__init__(layers)
Initialize a new MMCRProjectionHead instance.

        Args:
            input_dim:
                Number of input dimensions.
            hidden_dim:
                Number of hidden dimensions.
            output_dim:
                Number of output dimensions.
            num_layers:
                Number of hidden layers.
            batch_norm:
                Whether or not to use batch norms.
            use_bias:
                Whether or not to use bias in the linear layers.
        
lightly/models/modules/heads.py