The original datasets size is  [103x11x11] or [103x15x15] same datasets just different image sizes, it has 9 classes.
In “Train()” you can see the size of both tensors "output [400 9] , after FC.layer  " & “target [32 1] which is [batchSz 1]”

001.png1729×169 31.4 KB

The mismatching size error occurs in the highlighted line while trying to reshape both outputs, target into the same shape  :

002.png1770×395 62.7 KB

So what I wanna do exactly is finding a way in which flattening step before the FC should be of the form  [batchSz -1]

        self.ap = nn.AvgPool2d(2)
        self.fc1 = nn.Linear('***', 1024)  '***: i don't know what to set as an input here '
        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(1024, 9 )

    def forward(self, x):
        x = self.model(x)
        x = self.ap(x)
        x = x.view("**", -1)  " **: i wanna make it [batch_size -1] "
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x