src/transformers/models/levit/modeling_levit.py/MLPLayerWithBN/forward
class MLPLayerWithBN:    def forward(self, hidden_state):
        hidden_state = self.linear(hidden_state)
        hidden_state = self.batch_norm(hidden_state.flatten(0, 1)).reshape_as(hidden_state)
        return hidden_state