src/transformers/models/van/modeling_van.py/VanStage/forward
class VanStage:    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        hidden_state = self.embeddings(hidden_state)
        hidden_state = self.layers(hidden_state)
        # rearrange b c h w -> b (h w) c
        batch_size, hidden_size, height, width = hidden_state.shape
        hidden_state = hidden_state.flatten(2).transpose(1, 2)
        hidden_state = self.normalization(hidden_state)
        # rearrange  b (h w) c- > b c h w
        hidden_state = hidden_state.view(batch_size, height, width, hidden_size).permute(0, 3, 1, 2)
        return hidden_state