Visualize graph -- with tensor sizes please!


I’m currently using this package to plot graphs:

It’s good, but there’s something missing. I’d like my graphs to show output sizes of each operation. Maybe it’s already supposed to work. I guess some change in pytorch may have broken this functionality.

Looking at the code, you see that it traverses the graph using .grad_fn. It looks for tensors, and shows them in orange. But I never get any orange boxes in my graph.

My question: How can we get the tensors that were “saved for backward” while traversing the graph?