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.