Finally understood where I went wrong, just declaring nn.MaxPool2d(2) takes the kernel size as well as the stride as 2. I was expecting it to take the stride as 1 by default. So, in that case, the output size from the Max2d becomes 66. So 66*64 becomes 2304.