code: def get_norm_layer(
      num_features: int, num_splits: int, **kw: Any
  ) -> nn.Module:
      """Utility to switch between BatchNorm2d and SplitBatchNorm."""
      if num_splits > 0:
          return SplitBatchNorm(num_features, num_splits)
      else:
          return nn.BatchNorm2d(num_features)
  
file_path: lightly/lightly/models/batchnorm.py
imports: from __future__ import annotations
  from typing import Any
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 1 for calling a global function with example usage: get_norm_layer(num_features, num_splits) and returns: SplitBatchNorm, nn