src/transformers/models/jukebox/modeling_jukebox.py/JukeboxLayerNorm/forward
class JukeboxLayerNorm:    def forward(self, input):
        if input.numel() > self.max_numel:
            return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps).type_as(input)
        else:
            return super().forward(input).type_as(input)