torchvista: Interactive Forward-Pass Graphs Inside PyTorch Notebooks

The gap torchvista fills: forward-pass structure you can click on

Printing a model gives you a class hierarchy. A stack trace gives you a failure point. Neither shows the values flowing between layers. torchvista targets that middle ground: it runs a forward pass and draws the resulting graph in the notebook, where nodes can be dragged, zoomed, collapsed and clicked for parameter and attribute details. The README lists four features that map directly onto that job: an interactive graph with drag and zoom, collapsible nodes for hierarchical modules, error-tolerant partial visualization when errors such as shape mismatches arise, and info popups on node click. The intended user is someone iterating on a model definition in a notebook and wanting to see where a tensor's shape stops matching expectations. The error-tolerant mode is the most interesting of the four, because it means a broken forward pass still produces a partial picture rather than nothing. That is a debugging posture, not an inspection posture, and it tells you the author expected users to run this on models that do not yet work.

What trace_model actually does with model and inputs

The public API is a single function, trace_model, which takes a torch.nn.Module instance and the input or tuple of inputs to forward through it. The README's example builds a two-layer LinearModel, calls torch.randn(2, 10) for inputs, optionally sets model.eval(), and passes both to trace_model. The README does not describe the tracing mechanism itself; it points to a Towards Data Science blog post by the author for the under-the-hood explanation, so the internals are documented outside the repository. What the parameter table does reveal is the control surface. forced_module_tracing_depth defaults to None, which the table describes as tracing only user-defined modules. Setting it to an integer sets a maximum depth of module internals to trace. collapse_modules_after_depth defaults to 1, so nested modules start expanded one level deep, and 0 collapses everything while leaving interactive expansion available. show_non_gradient_nodes defaults to True and controls whether constants and values outside the gradient graph appear at all. Those four parameters are where the real tuning happens: how deep to descend, what to show initially, and whether non-gradient values clutter the canvas.

Getting it into a notebook: install, trace, export

Two install paths are documented. The pip route is pip install torchvista. The conda route is conda install -c conda-forge torchvista, and the README points at the torchvista-feedstock repository for details. The README is explicit that this runs from a web-based notebook: Jupyter, Colab, VS Code notebook and similar. The minimal call is two lines after your imports: from torchvista import trace_model, then trace_model(model, inputs). Everything else has a default. For file output, export_format accepts png, svg or html, and when it is omitted the graph is shown inside the notebook instead. export_path is where the documentation gets restrictive: only HTML format is currently supported with custom export paths, and a bare file name is created in the present working directory. So a png or svg export with a chosen destination is not a supported combination as documented. Layout is controlled by height (default 800 pixels) and width (pixels or a percentage, defaulting to full available width). show_module_attr_names defaults to False and switches node labels from class names to attribute names where available. There is also a test setup: pip install -e ".[test]" followed by pytest, with tests under tests/.

The compressed view is experimental, and the README says so

show_compressed_view defaults to False and is marked Experimental in the parameter table. It collapses repeating nodes of the same type with identical input and output dimensions into single repeat blocks. Two constraints are stated plainly. First, it currently only recognises repeating nodes within Sequential and ModuleList, so hand-rolled loops that build layers dynamically will not be compressed. Second, the README carries a warning that the feature might be expensive on large models. That is an honest disclosure and it should shape how you try it: enable it on a small model first, not on the largest thing in your repository. The same caution applies to forced_module_tracing_depth. Because None means only user-defined modules are traced, you get a readable graph by default, but the moment you set a numeric depth you are asking the tool to descend into library internals, and the cost of that descent is not quantified anywhere in the supplied material. Treat both parameters as opt-in experiments rather than defaults to tune casually.

Where torchvista is the wrong tool

Three boundaries are visible in the documentation. The first is environment: the README says web-based notebooks, naming Jupyter, Colab, Kaggle and VS Code notebooks. A terminal-only workflow or a non-web notebook front end is outside the stated scope, and nothing in the material suggests a fallback renderer. The second is scope of analysis. This is a forward-pass visualization. The parameter table mentions gradient nodes only in the negative, via show_non_gradient_nodes, which filters out constants and values outside the gradient graph. There is no mention of backward passes, memory accounting, kernel timings or operator-level profiling. If your question is why a step is slow, torchvista does not answer it. The third is export fidelity. Since export_path is documented as HTML-only, a pipeline that expects a PNG or SVG written to a chosen directory will need to render in the notebook and capture the output some other way. None of these are defects; they are the edges of a tool that does one thing.

How it differs from Netron and torchviz

The closest comparison in this space is Netron, which opens an exported model file and renders the graph in a standalone viewer. The difference in approach is where the graph comes from. Netron reads a serialized artifact, so it shows the graph the file format encodes. torchvista executes a forward pass with real inputs and draws what that execution produced, which is why it can show a partial graph when a shape mismatch throws. That also means torchvista needs a working Python process with the model instantiated and inputs available, while Netron needs a file. torchviz, the older graphviz-based helper, produces a static image of the autograd graph. torchvista's README emphasizes interaction (drag, zoom, collapse, click for parameter info) and multiple export formats, where a torchviz output is a fixed picture. If you want to inspect a model you received but cannot run, Netron fits. If you want to watch your own model's forward pass while you are editing it, torchvista's execution-based approach is the point.

Maintenance, releases and the GPL-3.0 question

The repository is not archived, the default branch is main, and the last push recorded is 2026-09-05, the same day as the v0.2.13 release. The release cadence visible in the material is three releases across roughly seven months: v0.2.11 in February 2026, v0.2.12 in May 2026, v0.2.13 in September 2026. That is a maintained but not fast-moving project, and the version numbers staying in 0.2.x tells you the author has not declared a stable API. Pin your version if you script against trace_model's signature. The licence is GPL-3.0. That is a copyleft licence, and it is worth flagging because visualization helpers are often assumed to be permissively licensed. If torchvista is imported into a notebook you keep to yourself, the practical effect is limited. If you distribute a product that bundles it, GPL-3.0 imposes obligations that a permissive licence would not. This is not legal advice; check your own distribution model with someone qualified. The upgrade cost itself is low, since the dependency surface appears to be PyTorch plus the notebook front end, and the test suite runs with pip install -e ".[test]" and pytest.

Editorial conclusion

Adopt torchvista if you debug nn.Module hierarchies and shape mismatches inside a Jupyter, Colab, Kaggle or VS Code notebook, and if a GPL-3.0 dependency is acceptable for your distribution model. Skip it if you need a profiler with per-operator timings, if your target is a non-web notebook front end, or if you cannot take a copyleft licence into a closed product. Before committing, verify three things in your own environment: that trace_model runs on your model with the default forced_module_tracing_depth=None, that the experimental show_compressed_view does not stall on your largest Sequential or ModuleList, and that your export path ends in .html, since export_path is documented as HTML-only.

Official sources

1. License: GPL-3.0
2. Project website
3. README
4. Releases
5. sachinhosmani/torchvista on GitHub