src/transformers/models/levit/modeling_levit.py/LevitSubsample/forward
class LevitSubsample:    def forward(self, hidden_state):
        batch_size, _, channels = hidden_state.shape
        hidden_state = hidden_state.view(batch_size, self.resolution, self.resolution, channels)[
            :, :: self.stride, :: self.stride
        ].reshape(batch_size, -1, channels)
        return hidden_state