code: def _maybe_expand_batchnorm_weights(
      model_dict, state_dict, num_splits
  ):
      """Expands the weights of the BatchNorm2d to the size of SplitBatchNorm."""
      running_mean = "running_mean"
      running_var = "running_var"
  
      for key, item in model_dict.items():
          # not batchnorm -> continue
          if not running_mean in key and not running_var in key:
              continue
  
          state = state_dict.get(key, None)
          # not in dict -> continue
          if state is None:
              continue
          # same shape -> continue
          if item.shape == state.shape:
              continue
  
          # found running mean or running var with different shapes
          state_dict[key] = state.repeat(num_splits)
  
      return state_dict
  
imports: import os
  import hydra
  import torch
  from hydra import utils
  from torch import nn
  from lightly.cli._cli_simclr import _SimCLR
  from lightly.embedding import SelfSupervisedEmbedding
  from lightly.models import ZOO
  from lightly.models import ResNetGenerator
  from lightly.models.batchnorm import get_norm_layer
  from lightly.utils.version_compare import version_compare
label_desc: Use image node_id 7 for calling a global function with example usage: _maybe_expand_batchnorm_weights(model_dict, state_dict, num_splits) and returns: state_dict