Solved by ptrblck in post #2 
              
              
                You could add padding=1 to the conv layers in b2 and b3, which would keep the spatial size. 
EDIT: For b3 you would need to pad the output in dim2, as you won’t get the same output shape for a kernel_size of 4 and padding of 1.