| 16 | |
| 17 | |
| 18 | class ChannelLayerNorm32(LayerNorm32): |
| 19 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 20 | DIM = x.dim() |
| 21 | x = x.permute(0, *range(2, DIM), 1).contiguous() |
| 22 | x = super().forward(x) |
| 23 | x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous() |
| 24 | return x |
| 25 |