I am trying to train the following model, but size mismatch error occurs in the classifier chunk of code.

Any help?

I am trying to train the following model, but size mismatch error occurs in the classifier chunk of code.

Any help?

Linear layer consists of matrix multiplication. For matrix multiplication, the number of columns of matrix 1 should be equal to number of rows of matrix 2. In your case the first one has [10 * 3584] while the second one has [25088 * 10] hence you get a mismatch. The 3584 and 25088 should match. I would print the output of avgpool and set the number of nodes in the linear layer according to that.