RuntimeError: The size of tensor a (8) must match the size of tensor b (2) at non-singleton dimension 0 with BCEWithLogitsLoss or BCELoss

Hi all,
I have a problem with the bce loss calculation.
Here is a small snippet of the code:

import torch.nn as nn
import torch
input = torch.tensor([[ 0.2363,  0.1395],[-0.0900, -0.0216],[ 0.4581,  0.1376],[ 0.0994,  0.1605],[-0.0092, -0.1196],[ 0.0835, -0.0903],[-0.0234, -0.1030],[ 0.0370, -0.1433]], device='cuda', requires_grad=True)
target = torch.tensor([0, 0, 1, 1, 0, 0, 0, 1], device='cuda')
weights = [1,1]
loss_function = nn.BCEWithLogitsLoss(weight=torch.tensor(tuple(weights), dtype=torch.float)).to('cuda')
loss = loss_function(input.argmax(1).float(), target.float())

and then I got this error message:

RuntimeError: The size of tensor a (8) must match the size of tensor b (2) at non-singleton dimension 0

I have tried to solve it, but I keep getting different errors every time.
Do you know how to work this out?

Thanks,
Roberto

You need to pass in the raw logits (don’t take argmax) and apply one_hot encoding to your targets, e.g. the shape for both the input and the target should be [batch_size, num_classes].

import torch.nn as nn
import torch
input = torch.tensor([[ 0.2363,  0.1395],[-0.0900, -0.0216],[ 0.4581,  0.1376],[ 0.0994,  0.1605],[-0.0092, -0.1196],[ 0.0835, -0.0903],[-0.0234, -0.1030],[ 0.0370, -0.1433]], device='cpu', requires_grad=True)
target = torch.tensor([0, 0, 1, 1, 0, 0, 0, 1], device='cpu')
target = nn.functional.one_hot(target)
weights = [1,1]
loss_function = nn.BCEWithLogitsLoss(weight=torch.tensor(tuple(weights), dtype=torch.float))
loss = loss_function(input, target.float())

Ok it works, thank you so much @soulitzer !