Hey guys!
I’m trying to train a neural network for multiclass semantic segmentation (FCN) using the CrossEntropyLoss function, but the following error occurs:
Preds: torch.Size([1, 21, 160, 240])
Target: torch.Size([1, 3, 160, 240])
The size of tensor a (21) must match the size of tensor b (3) at non-singleton dimension 1
As I am using a dataset with 21 classes (RGB 8 bit), my input should be correct. But the dimension of my target is different, this is a little confusing, I don’t know how to solve it.
Should I do some kind of transformation?
Any help or guidance on this will be greatly appreciated!
Part of the code where the error occurs:
def check_accuracy(loader, model, device="cuda"):
    num_correct = 0
    num_pixels = 0
    dice_score = 0
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            y = y.to(device).squeeze(1)
            preds = nn.Softmax(dim=1)(model(x))
            preds = (preds > 0.5).float()
            y = y.permute(0, 3, 1, 2) 
            print(f'Preds: {preds.shape}')
            print(f'Target: {y.shape}')
            num_correct += (preds == y).sum()         ####ERROR#####
            num_pixels += torch.numel(preds)
            dice_score += (2 * (preds * y).sum()) / (
                (preds + y).sum() + 1e-8
            )

    print(
        f"{num_correct}/{num_pixels} accuracy {num_correct/num_pixels*100:.2f}"
    )
    print(f"Dice Score : {dice_score/len(loader)}")
    model.train()