I am training a BERT based model with AutoModelForSequenceClassification. I want to use class weights for the loss function so I followed this discussion and implemented a BCEWithLogitsLoss.
Problem is, HuggingFace advices using num_labels=2 for Binary Classification with the Auto Class. When I try to implement BCEWithLogitsLoss I run into the error that
Value Error: Target size (torch.Size([32])) must be the same as input size (torch.Size([32, 2]))
I understand why this is happening, the size of the b_labels is the ground truth while the logits have the shape [32, 2] but I can’t figure out how to fix it. Here’s the code:
criterion = nn.BCEWithLogitsLoss(weight=class_weights, reduction='mean')
logits = outputs.logits
loss = criterion(logits, b_labels)
Here, class_weights is tensor([1.0712, 0.9377]) Should I just use CrossEntropyLoss instead? I have dataset imbalance that’s why I want to switch to custom loss and see what happens. Please guide me towards a solution.