Hi,
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)
model_graph.visual_graph
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.