title: Need a vectorized solution in pytorch
body: <p>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.</p>
  
  <p>Objective: To obtain patches of x of patch_size(given as argument).</p>
  
  <p>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. </p>
  
  <p>Used: PyTorch 0.4.1, (12 GB) Nvidia TitanX GPU.  </p>
  
  <p>The following is my implementation using two for loops</p>
  
  <pre><code>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
  </code></pre>
  
  <p>Thanks for your help in advance.</p>