Solved by iacob in post #16 
              
              
                For a given input shape, you can use the torchinfo (formerly torchsummary) package: 

Torchinfo provides information complementary to what is provided by print(your_model) in PyTorch, similar to Tensorflow’s model.summary()… 

Example: 
from torchinfo import summary

model = ConvNet()
batch_size = 1…