Is there a tool to visualize dataflow through a model?

I’m wondering if anyone has created a tool that will create some kind of dataflow graph of a PyTorch model showing the sizes of data at various points in the model. Even something that generates graphviz would be helpful.

Hmm there are couple of them. You can try tensorboard X which is a pytorch version of tensorboard. There is another one which is nice and group layers not to have a huge graph. The rest of them are very simple and graphs are usually illegible.