src/transformers/models/efficientformer/modeling_efficientformer.py/EfficientFormerFlat/forward
class EfficientFormerFlat:    def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor]:
        hidden_states = hidden_states.flatten(2).transpose(1, 2)
        return hidden_states