I see the issues. A cleaner way would be to set the in_features as the shape calculations, e.g. in_features=512*7*7. So that it would be clear the last conv layer outputs 512 channels and a volume of spatial size 7x7.



Utility function for calculating the shape of a conv output


I almost feel like it would be nice if this was just a method of the conv module: pass an input shape, and it returns the formula above.


This can be easily done by passing a random input to the conv layer:
nn.Conv2d(3, 6, 3, 1, 1)(torch.randn(1, 3, 24, 24)).shape

However, you would still need to know the input shape of your image.
I’m not sure there is a clean way to achieve the shape inference currently.
Maybe someone else has a good idea.