Pytorch CrossEntropy Loss, getting error: "RuntimeError: Boolean value of Tensor with more than one value is ambiguous"

Question

I have a classification model, producing predictions for 4 classes in a tensor of shape (256, 1, 4)...256 is the batch size, while the "1" for the second dimension is due to some model internal logic and can be removed:

preds.shape
torch.Size([256, 1, 4])

The corresponding annotations are one-hot encoded, in a tensor of the same shape:

targets.shape
torch.Size([256, 1, 4])

so, in every row there is only one non-zero element:

targets[0][0] = [1, 0, 0, 0]

I need to calculate the CrossEntropy loss of the model. I know that The CrossEntropyLoss expects class indices as the target, so I tried using argmax to determine the position of the 1 for each sample, and squeezed the extra dimension:

vpredictions_squeezed = cls_scores.squeeze(1)
targets = torch.argmax(targets.squeeze(1), dim=1)
losses = torch.nn.CrossEntropyLoss(predictions_squeezed, targets)

But I'm getting the error:

RuntimeError: Boolean value of Tensor with more than one value is ambiguous

What am I doing wrong here?

Answer 1

you can use direct call cross_entropy from torch.nn.functional

import torch.nn.functional as F
F.cross_entropy(predictions_squeezed, targets)

or you can rewrite your code, because this's class not a function:

loss = nn.CrossEntropyLoss()
output = loss(input, target)