RuntimeError: Assertion `input0 == target0 && input2 == target1 && input3 == target2' failed. size mismatch (got input: 1x12x174x4, target: 1x708x4)

Hi, I have the below CNN code, but I get an error when computing the Cross Entropy loss. Seems the shape of my y_hat is different from y_train, so how do I make the dimensions match?

X_train: torch.Size([1, 3, 708, 256])
Y_train: torch.Size([1, 708, 4])

class CNN(torch.nn.Module):
    def __init__(self):
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 24, kernel_size=5, stride=1),  
            nn.MaxPool2d(kernel_size=2, stride=0))
        self.conv2 = nn.Sequential(
            nn.Conv2d(24, 12, kernel_size=5, stride=1),
            nn.MaxPool2d(kernel_size=2, stride = 0))
        #self.drop_out = nn.Dropout()

        self.fc1 = nn.Linear(61, 12) 
        self.fc2 = nn.Linear(12, 4)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        #out = out.view(out.size(0), -1)
        #out = out.reshape(out.size(0), -1)
        #out = self.drop_out(out)
        out = self.fc1(out)
        out = self.fc2(out)
        return out
# defining the model
model = CNN()
model = model.float()
# defining the optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.07)
# defining the loss function
criterion = CrossEntropyLoss()
# checking if GPU is available
if torch.cuda.is_available():
    model = model.cuda()
    criterion = criterion.cuda()
def train(epoch):
    tr_loss = 0
    # getting the training set
    x_train, ytr= Variable(X_train), Variable(y_train)
    ytrain =
    # getting the validation set
    #x_val, y_val = Variable(val_x), Variable(val_y)
    # converting the data into GPU format

    # clearing the Gradients of the model parameters
    # prediction for training and validation set
    output_train = model(x_train.float())
   # output_train = model(x_train)
    print('Outputtrain shape',output_train.shape)

    # computing the training and validation loss
    loss_train = criterion(output_train, ytrain)


    # computing the updated weights of all the model parameters
    tr_loss = loss_train.item()

This part returns the error:

# defining the number of epochs
n_epochs = 5
# empty list to store training losses
train_losses = []
# empty list to store validation losses
val_losses = []
# training the model
for epoch in range(n_epochs):

Could you explain the shape of your target a bit?
Based on your model definition it looks like you are dealing with a multi-class classification, where you try to classify one out of four possible classes for each sample.

If so, the target should be a LongTensor with the shape [batch_size] containing class indices in the range [0, 3].

dim2 of your target looks as if it’s one-hot encoded, which won’t work with nn.CrossEntropyLoss. I’m not sure, what dim1 represents.

Also, it seems you’ve commented out the view operation, which flattens the output activation of your conv layer, and which would be the usual work flow.

A minor issue: currently you are not using any non-linearity between the last linear layers, which would basically create a single linear layer (with redundant parameters), so I would recommend to add an activation there. :wink:

Thanks for the detailed answer. I have updated the dimensions to be 1-D!