espnet2/torch_utils/get_layer_from_string.py

"""Minimal stub of espnet2.torch_utils.get_layer_from_string.get_layer

Only the activation functions used by TFGridNet are implemented.
"""
import torch.nn as nn


def get_layer(type_str: str):
    if type_str == "prelu":
        return nn.PReLU
    if type_str == "relu":
        return nn.ReLU
    if type_str == "gelu":
        return nn.GELU
    if type_str == "elu":
        return nn.ELU
    if type_str == "swish":
        return nn.SiLU
    if type_str == "sigmoid":
        return nn.Sigmoid
    if type_str == "tanh":
        return nn.Tanh
    raise ValueError(f"Unknown activation type: {type_str}")