import torch.nn as nn


def activation_helper(activation, dim=None):
    if activation == "sigmoid":
        act = nn.Sigmoid()
    elif activation == "tanh":
        act = nn.Tanh()
    elif activation == "relu":
        act = nn.ReLU()
    elif activation == "leakyrelu":
        act = nn.LeakyReLU()
    elif activation is None:

        def act(x):
            return x

    else:
        raise ValueError("unsupported activation: %s" % activation)
    return act
