src/transformers/models/cvt/modeling_cvt.py/CvtStage/forward
class CvtStage:    def forward(self, hidden_state):
        cls_token = None
        hidden_state = self.embedding(hidden_state)
        batch_size, num_channels, height, width = hidden_state.shape
        # rearrange b c h w -> b (h w) c"
        hidden_state = hidden_state.view(batch_size, num_channels, height * width).permute(0, 2, 1)
        if self.config.cls_token[self.stage]:
            cls_token = self.cls_token.expand(batch_size, -1, -1)
            hidden_state = torch.cat((cls_token, hidden_state), dim=1)

        for layer in self.layers:
            layer_outputs = layer(hidden_state, height, width)
            hidden_state = layer_outputs

        if self.config.cls_token[self.stage]:
            cls_token, hidden_state = torch.split(hidden_state, [1, height * width], 1)
        hidden_state = hidden_state.permute(0, 2, 1).view(batch_size, num_channels, height, width)
        return hidden_state, cls_token