| 57 | return self.norm(x) |
| 58 | |
| 59 | class LayerNorm(pl.LightningModule): |
| 60 | def __init__(self, features, eps=1e-6): |
| 61 | super().__init__() |
| 62 | self.a_2 = nn.Parameter(torch.ones(features)) |
| 63 | self.b_2 = nn.Parameter(torch.zeros(features)) |
| 64 | self.eps = eps |
| 65 | |
| 66 | def forward(self, x): |
| 67 | mean = x.mean(-1, keepdim=True) |
| 68 | std = x.std(-1, keepdim=True) |
| 69 | return self.a_2 * (x - mean) / (std + self.eps) + self.b_2 |
| 70 | |
| 71 | class SublayerConnection(pl.LightningModule): |
| 72 | def __init__(self, size): |