mat1 and mat2 shapes cannot be multiplied Pytorch lightning CNN
I'm working on a CNN for a project using Pytorch lightning. I don't know why am I getting this error. I've check the size of the output from the last maxpool layer and it is (-1,10,128,128). The error is for the linear layer. Any help would be appreciated.
def __init__(self):    
    super().__init__()
    self.model = nn.Sequential(
        nn.Conv2d(3,6,4,padding=2),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Conv2d(6,10,4,padding=2),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Linear(10*128*128,240),
        nn.ReLU(),
        nn.Linear(in_features = 240,out_features=101),
        nn.ReLU()
    )

My error looks like this:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2560x128 and 163840x240)


You have to match the dimension by putting the view method between the feature extractor and the classifier.
And it would be better not to use the relu function in the last part.
Code:
import torch
import torch.nn as nn

class M(nn.Module):
    def __init__(self):
        super(M, self).__init__()
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(3,6,4,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(6,10,4,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(10*128*128,240),
            nn.ReLU(),
            nn.Linear(in_features = 240,out_features=101)
        )
    def forward(self, X):
        X = self.feature_extractor(X)
        X = X.view(X.size(0), -1)
        X = self.classifier(X)
        return X

model = M()
# batch size, channel size, height, width
X = torch.randn(128, 3, 512, 512)
print(model(X))

