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.
Thanks!!