    def test_window_partitionn(self):
        '''
        test the grid_partition
        Returns:None
        '''
        x = torch.randn(1, 3, 28, 28)
        out = window_partition(x)
        self.assertEqual(out.shape, (16, 7, 7, 3))
test the grid_partition
Returns:None
tests/unittests/models/max_vit/test_max_vit_utils.py