| 19 | |
| 20 | |
| 21 | class LayerNorm(nn.Module): |
| 22 | def __init__(self, channels, eps=1e-5): |
| 23 | super().__init__() |
| 24 | self.channels = channels |
| 25 | self.eps = eps |
| 26 | self.gamma = nn.Parameter(torch.ones(channels)) |
| 27 | self.beta = nn.Parameter(torch.zeros(channels)) |
| 28 | |
| 29 | def forward(self, x): |
| 30 | x = x.transpose(1, -1) |
| 31 | x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) |
| 32 | return x.transpose(1, -1) |
| 33 | |
| 34 | |
| 35 | class TextEncoder(nn.Module): |