How to write a code that prints the shape after passing through each CNN layer

class CNN(nn.Module):
def init(self):
super(CNN, self).init()
self.conv1 = nn.Sequential(
nn.Conv2d(
in_channels=1,
out_channels=32,
kernel_size=5,
stride=1,
padding=2,
),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
)
self.conv2 = nn.Sequential(
nn.Conv2d(32, 64, 5, 1, 2),
nn.ReLU(),
nn.MaxPool2d(2),
)
# fully connected layer, output 10 classes
self.out = nn.Linear(64 * 7 * 7, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
# flatten the output of conv2 to (batch_size, 32 * 7 * 7)
x = x.view(x.size(0), -1)
output = self.out(x)
return output, x

The training data size is [ 60000, 28, 28 ].

My goal is to determine if the input volume equates to the output volume after passing through the convolutional neural layers. I know that to do this, I need to follow strictly the given equation, padding = (kernel_size - 1)/2. The kernel size can be any value, but I do not know whether the number of filters will affect the dimension. Do we need a softmax activation in the FC layer? What’s the effect of having it?

Hooks are a clean way to do this, but you could also simply print the shape after every conv call in your forward()