from torch.nn.modules.instancenorm import _InstanceNorm
from torch.nn.modules.batchnorm import _BatchNorm
from torch.nn import GroupNorm, LayerNorm


def assert_is_norm_layer(module) -> bool:
    """Check if the module is a norm layer.
    Args:
        module (nn.Module): The module to be checked.
    Returns:
        bool: Whether the module is a norm layer.
    """
    norm_layer_candidates = (_BatchNorm, _InstanceNorm, GroupNorm, LayerNorm)
    return isinstance(module, norm_layer_candidates)