How to plot models using torchviz or hiddenlayer

I am trying to plot models using torchviz and hiddenlayer but both gets errors.

torchviz - GitHub - waleedka/hiddenlayer: Neural network graphs and training metrics for PyTorch, Tensorflow, and Keras.
hiddenlayer - GitHub - szagoruyko/pytorchviz: A small package to create visualizations of PyTorch execution graphs

Common Code:

from transformers import AutoModel
model1 = AutoModel.from_pretrained("bert-base-uncased")

Code for torchviz :

from torchviz import make_dot, make_dot_from_trace


AttributeError: ‘BertModel’ object has no attribute ‘size’

Code for hiddenlayer:

import hiddenlayer as hl


AssertionError: Argument args must be provided for Pytorch models.

How to solve these errors or do you have any other model plot method?

There is a new package developed just for this purpose, called torchview. I tried your model with torchview with the following code

from transformers import AutoModel, AutoTokenizer
from torchview import draw_graph
model1 = AutoModel.from_pretrained("bert-base-uncased")

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

inputs = tokenizer("Hello world!", return_tensors="pt")

model_graph = draw_graph(model1, input_data=inputs)


The output is the following, (zoom in the image for clearity)

It also works very well with many other model types
The link for the project: torchview project
Disclaimer: I am the author of the package.

wow, thank you very much. I also try to build a library but I have no time so I skip it.