    def forward(self, x, size):
        """forward function"""
        hh, ww = size
        bb = x.shape[0]
        for layer in self.MHCA_layers:
            x = layer(x, (hh, ww))

        # return x"s shape : [bb, N, C] -> [bb, C, hh, ww]
        x = x.reshape(bb, hh, ww, -1).permute(0, 3, 1, 2).contiguous()
        return x
forward function
towhee/models/mpvit/mpvit.py