Solved by ptrblck in post #8 
              
              
                Add a print statement in your forward method so see the shape of x: 
    def forward(self, x):
        x = self.model(x)
        x = self.ap(x)
        x = x.view(x.size(0), -1)
        print(x.shape)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x…