RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x34400 and 41624x1376)

Hello! Im new using pytorch and I cant manage to find a proper solution to this error, could anybody help me?

class MaskCnnModel(ImageClassificationBase):
def init(self):
super().init()
self.network = nn.Sequential(
nn.Conv2d(3, 86, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(86, 172, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2, 2), # output: 172 x 43 x 43

        nn.Conv2d(172, 258, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.Conv2d(258, 258, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2, 2), # output: 344 x 21.5 x 21.5

        nn.Conv2d(258, 344, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.Conv2d(344, 344, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2, 2), # output: 688 x 10.8 x 10.8

        nn.Flatten(), 
        nn.Linear(344*11*11, 1376), 
        nn.ReLU(),
        nn.Linear(1376, 688),
        nn.ReLU(),
        nn.Linear(688, 2))
    
def forward(self, xb):
    return self.network(xb)

I think you need to replace

nn.Linear(344*11*11, 1376)

with

nn.Linear(344*10*10, 1376)

Note tha t in max pool the default is to have floored size (that is 21 instead of 21.5 and then 10 instead of 10.5). For example

>>> import torch
>>> layer = torch.nn.MaxPool2d(2, 2)
>>> a = torch.randn(1, 1, 5, 5)
>>> b = layer(a)
>>> b
tensor([[[[1.7285, 1.0654],
          [2.4715, 1.0278]]]])
>>> b.shape
torch.Size([1, 1, 2, 2])
>>>