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

Tupac: <p>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?</p>
<p>X_train: torch.Size([1, 3, 708, 256])<br>
Y_train: torch.Size([1, 708, 4])</p>
<pre><code class="lang-auto">class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 24, kernel_size=5, stride=1),  
            nn.Sigmoid(),
            nn.MaxPool2d(kernel_size=2, stride=0))
        self.conv2 = nn.Sequential(
            nn.Conv2d(24, 12, kernel_size=5, stride=1),
            nn.Sigmoid(),
            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
</code></pre>
<pre><code class="lang-auto"># 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()
    
print(model)
</code></pre>
<pre><code class="lang-auto">def train(epoch):
    model.train()
    tr_loss = 0
    # getting the training set
    x_train, ytr= Variable(X_train), Variable(y_train)
    ytrain = ytr.to(dtype=torch.int64)
    # getting the validation set
    #x_val, y_val = Variable(val_x), Variable(val_y)
    # converting the data into GPU format
    
    print('ytrain_shape',ytrain.shape)

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


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

    train_losses.append(loss_train)


    # computing the updated weights of all the model parameters
    loss_train.backward()
    optimizer.step()
    tr_loss = loss_train.item()
</code></pre>
<p>This part returns the error:</p>
<pre><code class="lang-auto"># 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):
    train(epoch)
</code></pre>

ptrblck: <p>Could you explain the shape of your target a bit?<br>
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.</p>
<p>If so, the target should be a <code>LongTensor</code> with the shape <code>[batch_size]</code> containing class indices in the range <code>[0, 3]</code>.</p>
<p>dim2 of your target looks as if it’s one-hot encoded, which won’t work with <code>nn.CrossEntropyLoss</code>. I’m not sure, what dim1 represents.</p>
<p>Also, it seems you’ve commented out the <code>view</code> operation, which flattens the output activation of your conv layer, and which would be the usual work flow.</p>
<p>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. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

Tupac: <p>Thanks for the detailed answer. I have updated the dimensions to be 1-D!</p>