code: def _get_byol_mlp(num_ftrs: int, hidden_dim: int, out_dim: int):
      """Returns a 2-layer MLP with batch norm on the hidden layer.
  
      Reference (12.03.2021)
      https://arxiv.org/abs/2006.07733
  
      """
      modules = [
          nn.Linear(num_ftrs, hidden_dim),
          nn.BatchNorm1d(hidden_dim),
          nn.ReLU(),
          nn.Linear(hidden_dim, out_dim),
      ]
      return nn.Sequential(*modules)
  
file_path: lightly/lightly/models/byol.py
imports: import warnings
  import torch
  import torch.nn
  from lightly.models._momentum import _MomentumEncoderMixin
  from lightly.models.modules import BYOLProjectionHead
label_desc: Use image node_id 1 for calling a global function with example usage: _get_byol_mlp(num_ftrs, hidden_dim, out_dim) and returns: nn