RuntimeError: Expected target size [128, 2], got [128]
The output and target shapes are torch.Size([128, 256, 2]) and torch.Size([128]). Trying a binary classification problem.
Appreciate any help on this.
Thank you,
The shapes are wrong as they don’t correspond to each other for a multi-class classification using nn.CrossEntropyLoss.
As described in the docs the output should have the shape [batch_size, nb_classes, *] containing logits and the target should have the shape [batch_size, *] containing class indices in the range [0, 1] or the same shape if probabilities are used.
In particular the model output should be permuted and you would then have to either reduce the dimemension with the size 256 or provide 256 targets for each sample.