Hi, till now I was using small model for MNIST like this -
Input tensor was (64x1x28x28) and I used this for flatten images to vector. So it was (64x784)
images = images.view(images.shape[0], -1)

input_size = 784
hidden_sizes = [128, 64]
output_size = 10

model = nn.Sequential(nn.Linear(input_size, hidden_sizes[0]),
                      nn.ReLU(),
                      nn.Linear(hidden_sizes[0], hidden_sizes[1]),
                      nn.ReLU(),
                      nn.Linear(hidden_sizes[1], output_size),
                      nn.LogSoftmax(dim=1))
print(model)

I changed model to this bigger model That starts with these layers.
How should I change input?
I have (64x1x28x28), how should i transofrm it for this model ?
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1)
        self.activation1 = nn.ReLU()

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.activation2 = nn.ReLU()
ETC.......
AT THE END I HAVE
       self.linear1 = nn.Linear(128 * 5 * 5, 10)