PyTorch Runtime Error: The size of tensor a (128) must match the size of tensor b (129) at non-singleton dimension 2

Hi,

I am a new PyTorch learner and got a problem recently in running the following code:

class CustomLoss(_Loss):
   
    def forward(self, original, disparity, mag):
        
        # original is the input
  
        loss=0
        target_overlap0=build_stack(original, 0, 1)  # target_overlap0: num x 1 x w x h
        target_overlap1=build_stack(original, 1, 1)  # target_overlap1: num x 1 x w x h
        target_overlap2=build_stack(original, 2, 1)  # target_overlap2: num x 1 x w x h
        target_overlap3=build_stack(original, 3, 1)  # target_overlap3: num x 1 x w x h
        target_mag1=build_stack(mag,1,1)
        target_mag2=build_stack(mag,2,1)
        target_mag3=build_stack(mag,3,1)

        
        w=original.size()[3]
        h=original.size()[2]  # size()[3] and size()[2] are respectively the width and height of the image
        disp=disparity[:,0,:,:] # take the first patch 
        # the second dimension is eliminated
        

        theta=torch.FloatTensor([[1,0,0],[0,1,0]])
        theta=theta.unsqueeze(0)
        theta=theta.repeat(original.size()[0],1,1)
#        original.size()[0] is the value of the patch size. 
#        '1' inside theta.repeat means it is exactly the same as the original argument.
#        If '2' inside the function instead, it means to copy the original argument once.
        
        theta=Variable(theta)  # A kind of structure. For calculating gradients more conveniently.
        grid_ori = F.affine_grid(theta, torch.Size((original.size()[0], 1, h, w)))
        # Affine transformation
        grid_ori=grid_ori.cuda(0)



        
        grid_shift0=torch.stack((-1*disp/w/2,            #5_5
                                -1*disp/h/2), dim=3)
        grid1 = grid_ori + grid_shift0
        warped_ori=F.grid_sample(target_overlap0, grid1, mode='bilinear')#, padding_mode='border')
        warped_ori_w=warped_ori*target_mag1
        target_overlap_w=target_overlap1*target_mag1
        diff1=F.mse_loss(warped_ori_w, target_overlap_w)
        
        

        
        grid_shift0=torch.stack((1*disp/w/2,             #4_4
                                -1*disp/h/2),dim=3)
        grid1=grid_ori+grid_shift0
        warped_ori=F.grid_sample(target_overlap0, grid1, mode='bilinear')#, padding_mode='border')
        warped_ori_w=warped_ori*target_mag2
        target_overlap_w=target_overlap2*target_mag2
        diff2=F.mse_loss(warped_ori_w, target_overlap_w)
        
        
#        grid_shift0=torch.stack((-1*disp/w/2,                #7_7
#                                1*disp/h/2),dim=3)
        
        grid_shift0=torch.stack((-1*disp/w/2,                #7_7
                                1*disp/h/2),dim=3)
        grid1=grid_ori+grid_shift0
        warped_ori=F.grid_sample(target_overlap0, grid1, mode='bilinear')#, padding_mode='border')
        warped_ori_w=warped_ori*target_mag3
        target_overlap_w=target_overlap3*target_mag3
        diff3=F.mse_loss(warped_ori_w, target_overlap_w)
        

#        grid_shift0=torch.stack((1*disp/w/2,                   #8_8
#                                 1*disp/h/2),dim=3)
#        grid1=grid_ori+grid_shift0
#        warped_ori=F.grid_sample(target_overlap0, grid1, mode='bilinear')#, padding_mode='border')
#        warped_ori_w=warped_ori*target_mag4
#        target_overlap_w=target_overlap4*target_mag4
#        diff4=F.mse_loss(warped_ori_w, target_overlap_w)
#        
        dx=gradx(disparity)[:,:,0:128,0:128]
        dy=grady(disparity)[:,:,0:128,0:128]
        
        TV=F.l1_loss(torch.abs(dx)+torch.abs(dy),0*disparity)
#        loss= diff1+diff2+diff3+diff4#+0.000002*TV
        loss= diff1+diff2+diff3+0.000002*TV

                
        return loss

Error occurred in this statement:

grid1 = grid_ori + grid_shift0

But I checked the dimension of grid_ori and grid_shift0. They are actually the same. I don’t know where this problem comes from. Is it a version problem?

Thank you all.