If you want to look at another example code, I have an old implementation of the architecture shown below.
The only difference is that I don’t use 1-max pooling but normal max-pooling. This implementation uses a nn.ModuleDict for a change :).
Please note that the code is rather verbose not make model very configurable, by specifying different conv_kernel_sizes as input parameter.