File size: 349 Bytes
8206b09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | from torch import nn
def get_activation(act_fn):
if act_fn in ["swish", "silu"]:
return nn.SiLU()
elif act_fn == "mish":
return nn.Mish()
elif act_fn == "gelu":
return nn.GELU()
elif act_fn == "relu":
return nn.ReLU()
else:
raise ValueError(f"Unsupported activation function: {act_fn}")
|