In general, the documentation of each operation like convolution:
https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
and maxpooling:
https://pytorch.org/docs/stable/generated/torch.nn.MaxPool2d.html
contains definitions for the shape calculations. However, using these still requires manual effort unless you write your own shape inference code.
In practice a simple hack that many use here is to simply print the shape before the linear layer runs, and to use the printed shape here. It’s a bit clunky, but afaict that is the simplest method at the moment. There may be better solutions in the future due to the move towards structured definitions that allow shape inference without actually running the code. rfcs/RFC-0005-structured-kernel-definitions.md at rfc-0005 · pytorch/rfcs · GitHub
As for both shapes working in your model; could you check that running the model actually works with both dimensions? I don’t think an error in dimensions would be caught until the model is actually run for the first time.