get activation layer Args: act_type (str): the activation type Returns: torch.nn.functional: the activation layer
(act_type)
| 262 | |
| 263 | |
| 264 | def get_activation_layer(act_type): |
| 265 | """get activation layer |
| 266 | |
| 267 | Args: |
| 268 | act_type (str): the activation type |
| 269 | |
| 270 | Returns: |
| 271 | torch.nn.functional: the activation layer |
| 272 | """ |
| 273 | if act_type == "gelu": |
| 274 | return lambda: nn.GELU() |
| 275 | elif act_type == "gelu_tanh": |
| 276 | return lambda: nn.GELU(approximate="tanh") |
| 277 | elif act_type == "relu": |
| 278 | return nn.ReLU |
| 279 | elif act_type == "silu": |
| 280 | return nn.SiLU |
| 281 | else: |
| 282 | raise ValueError(f"Unknown activation type: {act_type}") |
| 283 | |
| 284 | class IndividualTokenRefinerBlock(torch.nn.Module): |
| 285 | def __init__( |