class_docstr: Simulates multi-gpu behaviour of BatchNorm in one gpu by splitting.
  
  Implementation was adapted from:
  https://github.com/davidcpage/cifar10-fast/blob/master/torch_backend.py
  
  Attributes:
      num_features:
          Number of input features.
      num_splits:
          Number of splits.
code: def train(self, mode: bool = True) -> SplitBatchNorm:
      # lazily collate stats when we are going to use them
      if (self.training is True) and (mode is False):
          assert self.running_mean is not None
          self.running_mean = torch.mean(
              self.running_mean.view(
                  self.num_splits, self.num_features
              ),
              dim=0,
          ).repeat(self.num_splits)
          assert self.running_var is not None
          self.running_var = torch.mean(
              self.running_var.view(self.num_splits, self.num_features),
              dim=0,
          ).repeat(self.num_splits)
  
      return super().train(mode)
  
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 2 for calling the SplitBatchNorm obj's underlying member method code with example usage: obj.train(mode) and returns: super