Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back).
| 32 | |
| 33 | |
| 34 | class LayerNormFp32(nn.LayerNorm): |
| 35 | """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back).""" |
| 36 | |
| 37 | def forward(self, x: torch.Tensor): |
| 38 | orig_type = x.dtype |
| 39 | x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight, self.bias, self.eps) |
| 40 | return x.to(orig_type) |
nothing calls this directly
no outgoing calls
no test coverage detected