Here is the code for a text classification task I am doing. The issue seems to lie here. This is a multi class problem. I have 3 labels. I tried several things. I changed the format of the labels to integers and tried looking into the loss function. I am not sure what parameter needs to be changed.

def evaluate(model, dataloader_val):
    model.eval()
    model.train(False)
    
    loss_val_total = 0
    predictions, true_vals = [], []
    
    for batch in dataloader_val:
        
        batch = tuple(b.to(device) for b in batch)
        
        inputs = {'input_ids':      batch[0],
                  'attention_mask': batch[1],
                  'labels':         batch[2],
                 }

        with torch.no_grad():        
            outputs = model(**inputs)
            
        loss = outputs[0]
        
        logits = outputs[1]
        loss_val_total += loss.item()

        probs = torch.argmax(logits, dim = 1).detach().cpu().numpy()
        label_ids = inputs['labels'].cpu().numpy()
        predictions.append(probs)
        true_vals.append(label_ids)
        
    loss_val_avg = loss_val_total/len(dataloader_val) 
    
    predictions = np.concatenate(predictions, axis=0)
    true_vals = np.concatenate(true_vals, axis=0)
    
    ### after evaluating we resume model training
    model.train(True)
    
    return loss_val_avg, predictions, true_vals

And this is the error I get

ValueError                                Traceback (most recent call last)
<ipython-input-55-a095c6ad8f10> in <module>
     44                      }       
     45 
---> 46             outputs = model(**inputs)
ValueError: Expected input batch_size (2) to match target batch_size (4).