I’m not sure, if you would need to extract patches in your expected output is just the flattened tensor.
In my approach the first intermediate output would keep the patches as:
x = torch.arange(3* 8* 8).reshape((3,8,8))
out = x.unfold(2, 4, 4).unfold(1, 4, 4)
out
print(out.shape)
# torch.Size([3, 2, 2, 4, 4])
out = out.permute(1, 2, 0, 3, 4)
print(out)
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]],
           ...

If you want to get the flattened output wouldn’t a view operation just work?