I’m trying to build a CNN for classifying medical images into 5 classes. Here’s the architecture i used:
class ConvNet(nn.Module):
    
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.conv2 = nn.Conv2d(32, 32, 3)
        self.pool1 = nn.MaxPool2d(2, 2)

        self.conv3 = nn.Conv2d(32, 64, 3)
        self.conv4 = nn.Conv2d(64, 64, 3)
        self.pool2 = nn.MaxPool2d(2, 2)

        self.conv5 = nn.Conv2d(64, 128, 3)
        self.conv6 = nn.Conv2d(128, 128, 3)
        self.pool3 = nn.MaxPool2d(2, 2)

        self.conv7 = nn.Conv2d(128, 256, 3)
        self.conv8 = nn.Conv2d(256, 256, 3)
        self.pool4 = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(256*12*12,4096)
        self.fc2 = nn.Linear(4096, 5)
    
    def forward(self, input):
        x = self.pool1(self.conv2(F.relu(self.conv1(input))))
        x = self.pool2(self.conv4(F.relu(self.conv3(x))))
        x = self.pool3(self.conv6(F.relu(self.conv5(x))))
        x = self.pool4(self.conv8(F.relu(self.conv7(x))))
        #flatten the tensor for the FC
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

the input from trainloader is (10,3,256,256) and labels are of size “torch.Size([10])”.
and here is a snippet from the main:
for idx, batch in enumerate(trainloader):
        data = batch['image']
        labels = batch['label']
        if use_cuda:
            data = data.cuda()
            labels = labels.cuda()
        if labels.data.size()[0] == batch_size:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, labels)
            loss.backward()
            optimizer.step()

this line loss = criterion(output, labels.data) gives me error that error because they are not the same dimensions: input and target shapes do not match: input [10 x 5], target [10 x 1]
according to my understanding to cnn and the softmax layer at the end, the output should be 10x1(which is the highest class’s probability among those five ) but i found it to be 10x5. how can i fix this?
Thanks a lot