Deep learning with PyTorch:Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 2-dimensional input of size [32, 1728] instead
I am using deep learning with PyTorch to do some image classification. Whenever I try to train my model, the function forward fails. Could somebody explain me what are the reasons why the input size is wrong and how to fix this, please?
This is the code for my model, as well as my training loss and my optimizer:
model.fc = nn.Sequential(
        nn.Conv2d(1, 6, 9, padding=0),  # 64040
        nn.ReLU(),  #
        nn.AvgPool2d(2, stride=2),  # max 62020
        nn.Conv2d(6, 16, 11, padding=0),  # 161010
        nn.ReLU(),  # 161010
        nn.AvgPool2d(2, stride=2),  # 1655 = 400
        nn.Flatten(),
        nn.Linear(400, 200),
        nn.ReLU(),
        nn.Linear(200, 100),
        nn.ReLU(),
        nn.Linear(100, 3), 
        nn.LogSoftmax(dim=1))

criterion = nn.NLLLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.003)
model.to(device)

This is the validation function:
def validation(model, val_dataloader, criterion):

val_loss = 0
accuracy = 0

for images, labels in iter(val_dataloader):

    images, labels = images.to('cuda'), labels.to('cuda')

    output = model.forward(images)
    val_loss += criterion(output, labels).item()

    probabilities = torch.exp(output)
    
    equality = (labels.data == probabilities.max(dim=1)[1])
    accuracy += equality.type(torch.FloatTensor).mean()

return val_loss, accuracy

Finally, this is my training function:
def train_classifier():
  epochs = 10
  steps = 0
  print_every = 40

  model.to('cuda')

  for e in range(epochs):
  
      model.train()

      running_loss = 0

      for images, labels in iter(train_dataloader):
        images = images.view(images.shape[0], -1)  #this flattens it?
        
        steps += 1

        images, labels = images.to('cuda'), labels.to('cuda')

        optimizer.zero_grad()

        # training 
        output = model.forward(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
  
        if steps % print_every == 0:
        
            model.eval()
        
            # Turn off gradients for validation, saves memory and computations
            with torch.no_grad():
                validation_loss, accuracy = validation(model, validate_loader, criterion)
    
            print(&quot;Epoch: {}/{}.. &quot;.format(e+1, epochs),
                  &quot;Training Loss: {:.3f}.. &quot;.format(running_loss/print_every),
                  &quot;Validation Loss: {:.3f}.. &quot;.format(validation_loss/len(validate_loader)),
                  &quot;Validation Accuracy: {:.3f}&quot;.format(accuracy/len(validate_loader)))
    
            running_loss = 0
            model.train()
                    
train_classifier()

Error code:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-32-60a435d940e1&gt; in &lt;module&gt;()
     49             model.train()
     50 
---&gt; 51 train_classifier()

5 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    438                             _pair(0), self.dilation, self.groups)
    439         return F.conv2d(input, weight, bias, self.stride,
--&gt; 440                         self.padding, self.dilation, self.groups)
    441 
    442     def forward(self, input: Tensor) -&gt; Tensor:

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 2-dimensional input of size [32, 1728] instead

Any help is appreciated! Thank you!

Your network has both Conv2d layers as well as fully connected Linear layers - this is the source of the problem: Conv2d expects its inputs to be 4D: batch-channel-height-width. On the other hand, nn.Linear works of &quot;flattened&quot; features: batch-channel. Therefore, you need to &quot;flatten&quot; your data, but not before applying the network, but rather during its processing, where you have a Flattening layer.
