I wanted to announce that I have recently created a python package for visualizing pytorch models, called torchview.
If you wanted to checkout the repo for the project, see here
As far as I could have tried, it can capture any type of models with any type of operation in it.
For instances, the part of resnet18 graph from torchview (picture of the entire model is huge) is
One of the cool features is that it can return rolled version of recursive networks. For instance,
I added colab notebooks where you can explore some basic and complex/famous models.
I also listed todo list which I think should be fine to be implemented.
Any feedback/issue/pr is much appreciated!!!