I am using a transfer learning from a pretrained model. When I tried to do this:
loss_f = nn.CrossEntropyLoss(weight=torch.tensor([1,1,1])) output = model(volumes) labels = torch.FloatTensor(labels).long() #deleted .cuda() labels = labels.cuda() loss = loss_f(output, labels)
The error came up with:
ValueError: Expected target size (5, 50, 76, 80), got torch.Size([5, 2])
I printed the output and labels shape as follows:
output_shape = torch.Size([5, 2, 50 ,76, 80)] labels_shape = torch.Size([5, 2])
It seems that with the default documentation from CrossEntropyLoss is (N, C, D1, D2, …, Dk), which k is the dimensions of the image. I am trying to do binary classification, so the classes are 2. I am facing the problem with the shape differences between the output and the labels, could anyone comment on how to solve the problem.