def _prediction_mlp(in_dims: int, h_dims: int, out_dims: int) -> nn.Sequential:
    """Prediction MLP. The original paper's implementation has 2 layers, with
    BN applied to its hidden fc layers but no BN or ReLU on the output fc layer.

    Note that the hidden dimensions should be smaller than the input/output
    dimensions (bottleneck structure). The default implementation using a
    ResNet50 backbone has an input dimension of 2048, hidden dimension of 512,
    and output dimension of 2048

    Args:
        in_dims:
            Input dimension of the first linear layer.
        h_dims:
            Hidden dimension of all the fully connected layers (should be a
            bottleneck!)
        out_dims:
            Output Dimension of the final linear layer.

    Returns:
        nn.Sequential:
            The projection head.
    """
    l1 = nn.Sequential(
        nn.Linear(in_dims, h_dims), nn.BatchNorm1d(h_dims), nn.ReLU(inplace=True)
    )

    l2 = nn.Linear(h_dims, out_dims)

    prediction = nn.Sequential(l1, l2)
    return prediction
Prediction MLP. The original paper's implementation has 2 layers, with
    BN applied to its hidden fc layers but no BN or ReLU on the output fc layer.

    Note that the hidden dimensions should be smaller than the input/output
    dimensions (bottleneck structure). The default implementation using a
    ResNet50 backbone has an input dimension of 2048, hidden dimension of 512,
    and output dimension of 2048

    Args:
        in_dims:
            Input dimension of the first linear layer.
        h_dims:
            Hidden dimension of all the fully connected layers (should be a
            bottleneck!)
        out_dims:
            Output Dimension of the final linear layer.

    Returns:
        nn.Sequential:
            The projection head.
    
lightly/models/nnclr.py