Need a vectorized solution in pytorch
I'm doing an experiment using face images in PyTorch framework.  The input x is the given face image of size 5 * 5 (height * width) and there are 192 channels.

Objective: To obtain patches of x of patch_size(given as argument).

I have obtained the required result with the help of two for loops. But I want a better-vectorized solution so that the computation cost will be very less than using two for loops. 

Used: PyTorch 0.4.1, (12 GB) Nvidia TitanX GPU.  

The following is my implementation using two for loops

def extractpatches( x, patch_size): # x is bsx192x5x5
   patches = x.unfold( 2, patch_size ,  1).unfold(3,patch_size,1)
   bs,c,pi,pj, _, _  = patches.size()  #bs,192,
   cnt = 0
   p = torch.empty((bs,pi*pj,c,patch_size,patch_size)).to(device)
   s = torch.empty((bs,pi*pj, c*patch_size*patch_size)).to(device) 
  //Want a vectorized method instead of two for loops below
   for i in range(pi):
      for j in range(pj):
          p[:,cnt,:,:,:] = patches[:,:,i,j,:,:]
          s[:,cnt,:] = p[:,cnt,:,:,:].view(-1,c*patch_size*patch_size)
          cnt  = cnt+1
   return s


Thanks for your help in advance.

I think you can try this as following. I used some parts of your code for my experiment and it worked for me. Here l and f are the lists of tensor patches

l = [patches[:,:,int(i/pi),i%pi,:,:] for i in range(pi * pi)]
f = [l[i].contiguous().view(-1,c*patch_size*patch_size) for i in range(pi * pi)] 


You can verify the above code using toy input values.
Thanks.
