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)
Utility to switch between BatchNorm2d and SplitBatchNorm.
lightly/models/batchnorm.py