Subclass torch's LayerNorm to handle fp16.
| 151 | |
| 152 | |
| 153 | class LayerNorm(nn.LayerNorm): |
| 154 | """Subclass torch's LayerNorm to handle fp16.""" |
| 155 | |
| 156 | def forward(self, x: torch.Tensor): |
| 157 | orig_type = x.dtype |
| 158 | ret = super().forward(x.type(torch.float32)) |
| 159 | return ret.type(orig_type) |
| 160 | |
| 161 | |
| 162 | class QuickGELU(nn.Module): |