Function obtain_norm

code/utils.py:205–216

Return normalization function module

(name)

def obtain_norm(name):
    """
    Return normalization function module
    """
    if name == "layernorm":
        norm = nn.LayerNorm
    elif name == "batchnorm":
        norm = nn.BatchNorm1d
    else:
        raise NotImplementedError("{} is not implemented.".format(name))

    return norm


# ======================================================================