Thanks for the quick reply. Though it gets the shapes correct it sort of doesnt get the ordering of each sample correct.
Consider this code:
>>> x = torch.arange(3* 8* 8).reshape((3,8,8))
>>> p = 4
>>> res = []
>>> for i in range(1, 3):
...     for j in range(1,3):
...             res.append(x[:, p*(i-1):p*i, p*(j-1):p*j].flatten())
>>> res[0]
tensor([  0,   1,   2,   3,   8,   9,  10,  11,  16,  17,  18,  19,  24,  25,
         26,  27,  64,  65,  66,  67,  72,  73,  74,  75,  80,  81,  82,  83,
         88,  89,  90,  91, 128, 129, 130, 131, 136, 137, 138, 139, 144, 145,
        146, 147, 152, 153, 154, 155])

Whereas in yor case it looks like:
>>> out[0]
tensor([  0,   8,  16,  24,   1,   9,  17,  25,   2,  10,  18,  26,   3,  11,
         19,  27,  64,  72,  80,  88,  65,  73,  81,  89,  66,  74,  82,  90,
         67,  75,  83,  91, 128, 136, 144, 152, 129, 137, 145, 153, 130, 138,
        146, 154, 131, 139, 147, 155])

Any idea how to fix this?