PyTorch 3 reshaping error

I will assume your input images are probably of size 200x200px (by size I mean here height x width, not taking the number of channels into account).

While your nn.Conv2d layers are defined to output tensors of the same size (with 32 channels for conv1 and 64 channels for con2), the F.max_pool2d are defined in such a way they divide height and width by 2.

So after 2 max-pooling operations, your tensors are of size 200 / (2 * 2) x 200 / (2 * 2) = 50x50px. With the 64 channels from conv2, you get 64 * 50 * 50 = 160000 elements.

Now, you need to adapt your view() so that it converts those inputs of shape (batch_size, 64, 50, 50) into (batch_size, 64 * 50 * 50) (to preserve the number of elements). You need to similarly adapt your 1st fully-connected layer.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class Net(nn.Module):

    def __init__(self):
        super(Net,self).__init__()
        self.conv1  = nn.Conv2d(3,32,5,padding=2) # 1 input, 32 out, filter size = 5x5, 2 block outer padding
        self.conv2  = nn.Conv2d(32,64,5,padding=2) # 32 input, 64 out,  filter size = 5x5, 2 block padding
        self.fc1    = nn.Linear(64*50*50,1024) # Fully connected layer
        self.fc2    = nn.Linear(1024,2) #Fully connected layer 10 out.

    def forward(self,x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2) # Max pool over convolution with 2x2 pooling
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2) # Max pool over convolution with 2x2 pooling
        x = x.view(-1,64*50*50) # tensor.view() reshapes the tensor
        x = F.relu(self.fc1(x)) # Activation function after passing through fully connected layer
        x = F.dropout(x, training=True) #Dropout regularisation
        x = self.fc2(x) # Pass through final fully connected layer
        return F.log_softmax(x) # Give results using softmax

model = Net()
print(model)

x = np.ones((1, 3, 200, 200))
x = torch.tensor(x)
x = model.forward(x)
print(x)