Solved by ptrblck in post #2 
              
              
                You already have the results. The 4 patches are stored as 2x2 patches in dim2 and dim3, so you can just call .contiguous().view(batch_size, S, -1, size, size) on the result.