code: 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
  
file_path: lightly/lightly/models/nnclr.py
imports: import warnings
  import torch
  import torch.nn
  from lightly.models.modules import NNCLRPredictionHead, NNCLRProjectionHead
label_desc: Use image node_id 1 for calling a global function with example usage: _prediction_mlp(in_dims, h_dims, out_dims) and returns: prediction