code: def _projection_mlp(
      num_ftrs: int, h_dims: int, out_dim: int, num_layers: int = 3
  ) -> nn.Sequential:
      """Projection MLP. The original paper's implementation has 3 layers, with
      BN applied to its hidden fc layers but no ReLU on the output fc layer.
      The CIFAR-10 study used a MLP with only two layers.
  
      Args:
          in_dims:
              Input dimension of the first linear layer.
          h_dims:
              Hidden dimension of all the fully connected layers.
          out_dims:
              Output Dimension of the final linear layer.
          num_layers:
              Controls the number of layers; must be 2 or 3. Defaults to 3.
  
      Returns:
          nn.Sequential:
              The projection head.
      """
      l1 = nn.Sequential(
          nn.Linear(num_ftrs, h_dims),
          nn.BatchNorm1d(h_dims),
          nn.ReLU(inplace=True),
      )
  
      l2 = nn.Sequential(
          nn.Linear(h_dims, h_dims),
          nn.BatchNorm1d(h_dims),
          nn.ReLU(inplace=True),
      )
  
      l3 = nn.Sequential(
          nn.Linear(h_dims, out_dim), nn.BatchNorm1d(out_dim)
      )
  
      if num_layers == 3:
          projection = nn.Sequential(l1, l2, l3)
      elif num_layers == 2:
          projection = nn.Sequential(l1, l3)
      else:
          raise NotImplementedError(
              "Only MLPs with 2 and 3 layers are implemented."
          )
  
      return projection
  
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 2 for calling a global function with example usage: _projection_mlp(num_ftrs, h_dims, out_dim, num_layers) and returns: projection