Is there any function in pytorch like plot_model function in the tensor flow to print the network architecture?

I want to print the network architecture like the one shown in the figure. Is there any function in pytorch which can do so?

You can use torchviz library for the same!

This is my code which throwing errors. Can you provide me some more reference about the torchviz?

!pip install git+
from torchviz import make_dot

Take a look at this issue:

Another library is torchview, which is very close to plot_model of keras as it can capture module hierarchy. The similarity to plot_model API was a big factor in the design of the library
For instance, output for mlp model is the following

Disclaimer: I am the author of library