Julio_Marco_A_Silva:

As you know, BCEWithLogitsLoss accepts a vector of integers (one for each element in the batch) and I have a one-hot vector of two elements as the output of my network.


That’s not correct, as nn.BCEWithLogitsLoss expects raw logits as the model output (FloatTensors) and a target tensor with the same shape and type as the output containing values in [0, 1].
I’m a bit confused about your use case at the moment.
If you are dealing with a multi-label classification (each sample can belong to more than a single class), remove the softmax in your model and just pass the logits to the criterion without any max operations.
On the other hand, if you are dealing with a multi-class classification (each sample belongs to one class only), still remove the softmax, use nn.CrossEntropyLoss as the criterion, and pass the targets as class indices (e.g. by using torch.argmax(target).