so (excluding the batch dimension):
input: 1X28X28
after conv2d: 64X28X28
after maxpool: 64X13X13
after conv2d: 64X11X11
after maxpool: 64X5X5
after conv2d: 128X3X3
after maxpool: 128X1X1
now you have tensor of size (N, 128X1X1) which you view as (N, 128) by flattening it.
That’s why the number of input features to the first linear layer is 128