model = Sequential()

mnist_model.add(Conv2D(32, kernel_size=(3,3), activation='relu', input_shape=(28,28,1)))
mnist_model.add(MaxPool2D(pool_size=(2,2)))
mnist_model.add(Dropout(0.25))

mnist_model.add(Conv2D(64, kernel_size=(3,3), activation='relu'))
mnist_model.add(MaxPool2D(pool_size=(2,2)))
mnist_model.add(Dropout(0.25))

mnist_model.add(Flatten())
mnist_model.add(Dense(256, activation='relu'))
mnist_model.add(Dropout(0.5))
mnist_model.add(Dense(10, activation='softmax'))

What do I do in this case where one convolution layer is followed by another?
Below is my implementation. Is the architecture same as above same as that described by the above tensorflow code?
class NeuralNet(nn.Module):
  
    def __init__(self):
        super(NeuralNet, self).__init__()

        # spatial size: (28,28)    # number of channels = 1
        self.conv1 = nn.Conv2d(1, 32, kernel_size=(3,3))
        self.pool1 = nn.MaxPool2d(2)
        self.drop1 = nn.Dropout(0.25)

        # spatial size: (13, 13)    # number of channels = 32
        self.conv2 = nn.Conv2d(32, 32*64, kernel_size=(3,3))
        self.pool2 = nn.MaxPool2d(2)
        self.drop2 = nn.Dropout(0.25)

        # spatial size: (5,5)    # number of channels = 32*64 (applys 64 different filters on each channel?)
        # therefore number of nodes for the hidden layer 32*64*5*5
        self.linear3 = nn.Linear(32*64*5*5, 256)
        self.drop3 = nn.Dropout(0.5)

        self.linear4 = nn.Linear(256,10)

        self.relu = nn.ReLU()
        self.softmax = nn.Softmax()
        
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool1(x)
        x = self.drop1(x)
        
        x = self.relu(self.conv2(x))
        x = self.pool2(x)
        x = self.drop2(x)
        
        x = x.view(-1, self.num_flat_features(x))
        x = self.relu(self.linear3(x))
        x = self.drop3(x)
        
        x = x.view(-1, self.num_flat_features(x))
        x = self.linear4(x)
        
        return x
      
    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features
        
        
model = NeuralNet()

batch_size, C, H, W = 1, 1, 28, 28
x = torch.randn(batch_size, C, H, W)
output = model(x)