[Python] How to decide the 'input_size' parameter of torchsummary.summary(model=model.policy,...

This is my CNN network printed by 'print(model.policy)':

CnnPolicy(
(actor): Actor(
(features_extractor): CustomCNN(
(cnn): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
(1): ReLU()
(2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
(3): ReLU()
(4): Flatten(start_dim=1, end_dim=-1)
)
(linear): Sequential(
(0): Linear(in_features=6, out_features=128, bias=True)
(1): ReLU()
)
)
(mu): Sequential(
(0): Linear(in_features=128, out_features=256, bias=True)
(1): ReLU()
(2): Linear(in_features=256, out_features=128, bias=True)
(3): ReLU()
(4): Linear(in_features=128, out_features=3, bias=True)
(5): Tanh()
)
)


When I try to print the network architecture using torchsummary.summary(model=model.policy, input_size=(1, 32, 32)). I got the following error: RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x50176 and 6x128)

I have tried lots of 'input_size' combinations, but all were wrong.

I want to know how to choose the 'input-size' parameter?

Continue reading...