Get the normalization layer. Args: norm_layer (str): The type of normalization layer. Returns: norm_layer (nn.Module): The normalization layer.
(norm_layer)
| 244 | |
| 245 | |
| 246 | def get_norm_layer(norm_layer): |
| 247 | """ |
| 248 | Get the normalization layer. |
| 249 | |
| 250 | Args: |
| 251 | norm_layer (str): The type of normalization layer. |
| 252 | |
| 253 | Returns: |
| 254 | norm_layer (nn.Module): The normalization layer. |
| 255 | """ |
| 256 | if norm_layer == "layer": |
| 257 | return nn.LayerNorm |
| 258 | elif norm_layer == "rms": |
| 259 | return RMSNorm |
| 260 | else: |
| 261 | raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") |
| 262 | |
| 263 | |
| 264 | def get_activation_layer(act_type): |