src/transformers/models/esm/modeling_esmfold.py/EsmFoldLayerNorm/forward
class EsmFoldLayerNorm:    def forward(self, x):
        d = x.dtype
        if d is torch.bfloat16 and not is_deepspeed_initialized():
            with torch.cuda.amp.autocast(enabled=False):
                out = nn.functional.layer_norm(x, self.c_in, self.weight.to(dtype=d), self.bias.to(dtype=d), self.eps)
        else:
            out = nn.functional.layer_norm(x, self.c_in, self.weight, self.bias, self.eps)

        return out