but I get an error:
RuntimeError: The size of tensor a (352) must match the size of tensor b (2) at non-singleton dimension 3.

How can I do this? From what I’ve gathered, weights is expected to have the same dimensions as the batch. This is doing batch weighting, but what I actually want is class weighting. From what I understand, BCEWithLogits does have both weights and pos_weights, but I would like to avoid changing the criterion.

I guess there might be workarounds using nn.BCELoss with the sample weighting by computing the weight per batch using the class distribution, but I would rather switch to nn.BCEWithLogitsLoss allowing you to use pos_weights directly and for better numerical stability.

A simple choice for pos_weights would be num_negatives / num_positives as it will weight the positive samples so that they have approx. the same influence on the loss.

I had a doubt, I am using focal loss to tackle class imbalance. And I do not have a sigmoid layer at the end of my model. So I should be using BCEWithLogits correct?
Also, these class weights are obtained using the inverse class frequencies:
This is my focal loss implementation:
class FocalLoss(nn.Module):
def init(self, alpha=0.25, gamma=2.0, reduction=‘mean’):
super(FocalLoss, self).init()
self.alpha = alpha
self.gamma = gamma
self.reduction = reduction