Subclass torch's LayerNorm (with cast back to input dtype).
| 23 | |
| 24 | |
| 25 | class LayerNorm(nn.LayerNorm): |
| 26 | """Subclass torch's LayerNorm (with cast back to input dtype).""" |
| 27 | |
| 28 | def forward(self, x: torch.Tensor): |
| 29 | orig_type = x.dtype |
| 30 | x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) |
| 31 | return x.to(orig_type) |
| 32 | |
| 33 | |
| 34 | class LayerNormFp32(nn.LayerNorm): |
nothing calls this directly
no outgoing calls
no test coverage detected