Hi,
first, I personally think modifying the source code (well a copy) is a great way to do this!
typically, the size mismatch is when you move from the conv-layers (which don’t care about the height and width unless you have to few pixels left) to the Linear (aka fully connected) layers. These need to know the size of the inputs.
My not terribly elegant way of finding out is to put a print (x.size()) right before the first call to a linear layer, reading off the size, remove the print and adjust the network definition.
Best regards
Thomas