Your output formula is missing the dilation and also the subtraction from the kernel size.
The Conv2d docs 210 show you the formula which is used.
That being said, your printed conv layer block would keep the spatial dimensions equal in the first layers, since conv layers with a kernel size of 3 and padding of 1 would not reduce the height or width of the activation. The max pooling layer would halve the spatial dimensions.
Based on the out_channels=256 I thus assume that the input to the conv block would be [batch_size, 128, 8, 8], the output thus [batch_size, 256, 4, 4], which would be flattened to [batch_size, 256*4*4=4096].
The Convolution arithmetic tutorial 197 is a very good website to learn more about how convolution layers perform the window sliding.