def get_act()

in jcm/models/layers.py [0:0]


def get_act(config):
    """Get activation functions from the config file."""

    if config.model.nonlinearity.lower() == "elu":
        return nn.elu
    elif config.model.nonlinearity.lower() == "relu":
        return nn.relu
    elif config.model.nonlinearity.lower() == "lrelu":
        return functools.partial(nn.leaky_relu, negative_slope=0.2)
    elif config.model.nonlinearity.lower() == "swish":
        return nn.swish
    else:
        raise NotImplementedError("activation function does not exist!")