How to add Dense Layer after the last Conv layer when you don't know the shape of the last Conv layer activation map?

    layer1 = nn.Sequential(
        nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
        nn.BatchNorm2d(16),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2))
    layer2 = nn.Sequential(
        nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
        nn.BatchNorm2d(32),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2, stride=2))
    # How to add Dense Layer here ?

One simple method would be to add a print statement with the shape information after layer2 and define the number of input features based on this.
Otherwise, you could also calculate the shape manually using the convolution / pooling arguments, but the first approach might be faster and simpler.