Solved by ptrblck in post #2 
              
              
                tensor.unfold should yield the same result: 
a = torch.randn(64, 3, 16, 16)
kernel_size = 4
kernel_stride = 4
a = a.unfold(2, kernel_size, kernel_stride).unfold(3, kernel_size, kernel_stride)
a = a.contiguous().view(a.size(0), a.size(1), -1, a.size(4), a.size(5))
print(a.shape)
> torch.Size([64, 3, …