UserWarning: Using a target size (torch.Size([4])) that is different to the input size (torch.Size([1, 4]))

Hello, all
I found this Warning:

/home/pc/.local/lib/python3.8/site-packages/torch/nn/modules/loss.py:528: UserWarning: Using a target size (torch.Size([4])) that is different to the input size (torch.Size([1, 4])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.

My code is:

def load_filtered_state_dict(model, snapshot):
    # By user apaszke from discuss.pytorch.org
    model_dict = model.state_dict()
    snapshot = {k: v for k, v in snapshot.items() if k in model_dict}
    model_dict.update(snapshot)
    model.load_state_dict(model_dict)


if __name__ == '__main__':

    args = parse_args()
    cudnn.enabled = True
    num_epochs = args.num_epochs
    batch_size = args.batch_size
    gpu = args.gpu_id
    b_scheduler = args.scheduler
    dataset_name = args.dataset
    snapshot_name = args.snapshot
    # =====================learn_info tar ==================
    datetime_ = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    time_ = time.time()
    print("datetime_",datetime_,"time_",time_)

    if not os.path.exists('output/snapshots'):
        os.makedirs('output/snapshots') #Redhwan added exist_ok=True

    summary_name = '{}_{}_bs{}'.format(
        dataset_name, datetime_, args.batch_size)

    if not os.path.exists('output/snapshots/{}'.format(summary_name)):
        os.makedirs('output/snapshots/{}'.format(summary_name)) #Redhwan added exist_ok=True
    #=====================learn_info txt==================
    if not os.path.exists('output/learn_info'):
        os.makedirs('output/learn_info') #Redhwan added exist_ok=True

    name_txt = '{}_{}'.format(
        dataset_name, datetime_)

    if not os.path.exists('output/learn_info/{}'.format(name_txt)):
        os.makedirs('output/learn_info/{}'.format(name_txt)) #Redhwan added exist_ok=True


    backbone_name = 'RepVGG-B1g2'     
    backbone_file = 'RepVGG-B1g2-train.pth'


    model = RotationNet(backbone_name,
                       backbone_file,
                       deploy=False,
                       pretrained=True)
   
    if not args.snapshot == '':
        saved_state_dict = torch.load(args.snapshot)
        model.load_state_dict(saved_state_dict['model_state_dict'])

    print('Loading data.')

    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225])

    transformations = transforms.Compose([transforms.RandomResizedCrop(size=224,scale=(0.8,1)),
                                          transforms.ToTensor(),
                                          normalize])

    pose_dataset = datasets2.getDataset(
        args.dataset, args.data_dir, args.filename_list, transformations)
    print('pose_dataset_____________', pose_dataset)
    train_loader = torch.utils.data.DataLoader(
        dataset=pose_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4)

    model.cuda(gpu)
    
    crit = torch.nn.MSELoss().cuda(gpu)
    # softmax = nn.Softmax().cuda(gpu)
    optimizer = torch.optim.Adam(model.parameters(), args.lr) #Adam      
    print('optimizer', optimizer)


    if not args.snapshot == '':
        optimizer.load_state_dict(saved_state_dict['optimizer_state_dict'])

    #milestones = np.arange(num_epochs)
    milestones = [10, 20]
    scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, milestones=milestones, gamma=0.5)

    
    for epoch in range(num_epochs):
        loss_sum = .0
        iter = 0

        for i, (images, labels, cont_labels, name) in enumerate(train_loader):
            iter += 1
            images = torch.Tensor(images).cuda(gpu)

           
            pred_mat = model(images)
          

           
            print('pred_mat.shape', pred_mat.shape, 'labels.cuda(gpu)', labels.cuda(gpu).shape)
            loss = crit(labels.cuda(gpu), pred_mat)
          

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()    
            loss_sum += loss.item()                   
              
               

    outfile.close()

I am trying to train my code, and it is working, but it is Warning me.

When I check the target size and input size by this line as you see:

print('pred_mat.shape', pred_mat.shape, 'labels.cuda(gpu)', labels.cuda(gpu).shape)

The output is:

pred_mat.shape torch.Size([16, 4]) labels.cuda(gpu) torch.Size([16, 4])

What should I do in this case for this UserWarning?

or

How do I correctly ensure they have the same size?

Thank you in advance!

Based on the printed shapes in the warning message I guess the last batch might contain a single sample which might then be handled differently. Check if this could be the case and either drop it via drop_last=True in the DataLoader or by unsqueezing the apparently missing batch dimension in one of the tensors. You could e.g. check the tensor.dim() attribute and .unsqueeze a dimension if only a single dimension is returned.

Thank you so much for your help.
I increase the batch_size from 16 to 40. I did not see this warning.