Calculation for the input to the Fully Connected Layer

You are absolutely right, that is another possible solution!
Although, as in your post, the model contains a nn.MaxPool2d, you may want to use same pooling layer.
On top of that, before the Linear layer, the model is still using images so 2d pooling is needed.

So, The alternative is nn.AdaptiveMaxPool2d. See its documentation.

The benefit of using Adaptive pooling layer is that you explicitly define your desired output size so not matter what input size is, model always will produce tensors with the identical shape. Also, it is has been used in official PyTorch implementation of ResNet models right before Linear layer. Please see this post.

1 Like