szagoruyko / pytorchviz

A small package to create visualizations of PyTorch execution graphs
MIT License
3.24k stars 279 forks source link

Option to show number of parameters in each node #22

Open turian opened 6 years ago

turian commented 6 years ago

torchsummary has the capability to show the number of parameters in each node, and the shape of the tensor. Could you add this functionality too?