Hey, newbie here
Trying to build this classification model
image1685×700 165 KB
class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)  
    self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3) 
    self.pool = nn.MaxPool2d(2,2)
    self.fc1 = nn.Linear(32*7*7, 128) # or 64, not sure here
    self.fc2 = nn.Linear(128, 10)
    self.dropout = torch.nn.Dropout(p=0.5)
    self.relu = torch.nn.ReLU()

  def forward(self, x):
    x = self.conv1(x)
    x = self.relu(x)
    x = self.pool(x)
    x = self.conv2(x) 
    x = self.relu(x)
    x = self.pool(x) 
    x = x.reshape(x.size(0), -1) 
    x = self.fc1(x) 
    x = self.dropout(x)
    prediction = self.fc2(x) 

    return prediction

Outputs this error
Traceback (most recent call last):
  File "main.py", line 36, in <module>
    outputs = net(inputs)
  File "/home/kaneda/anaconda3/envs/home/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/kaneda/Documents/mnist_dataset/models/cnn.py", line 23, in forward
    x = self.fc1(x) # 128 * 10
  File "/home/kaneda/anaconda3/envs/home/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/kaneda/anaconda3/envs/home/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 91, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/kaneda/anaconda3/envs/home/lib/python3.8/site-packages/torch/nn/functional.py", line 1674, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: mat1 dim 1 must match mat2 dim 0

Any ideas? Using pytorch 1.6