import torch


def get_norm_layer(norm: str, channels: int):
    """
    Args:
        norm: str
            type of normalization: `layernorm`, `batchnorm`, `instancenorm`
        channels: int
            size of features for normalization
    """
    if norm == "layernorm":
        return torch.nn.LayerNorm(channels)
    elif norm == "batchnorm":
        return torch.nn.BatchNorm1d(channels)
    elif norm == "instancenorm":
        return torch.nn.InstanceNorm1d(channels)
    else:
        return torch.nn.Identity()