Well it’s 4*4*20 = 320.
You can calculate the shape for the forward pass for each operation.
While a convolution with kernel_size=5 and no padding shrinks the activation by 4 pixels in height and width, a max pooling of kernel_size=2 and stride=2 pools the activation to half its size.
From the forward pass:
# x.shape [1, 28, 28]
x = self.conv1(x) # x.shape: [10, 24, 24]
x = F.max_pool2d(x, 2) # x.shape: [10, 12, 12]
x = F.relu(x)
x = self.conv2(x) # x.shape: [20, 8, 8]
x = self.conv2_drop(x)
x = F.max_pool2d(x, 2) # x.shape: [20, 4, 4]
x = F.relu(x)
x = x.view(-1, 320)