Expand the architecture.
Start with the whole model. Expand a module to reveal its layers, operations, and skip connections. Collapse everything to return to the overview.
YOUR PYTORCH MODEL, MADE VISIBLE
A VS Code extension
Visualize your
PyTorch model.
Turn nested modules into a graph you can explore. Expand blocks into operations, follow tensor shapes, and jump from any layer to its Python source in VS Code.
Expand nested layersInspect runtime shapesReview code changes
EXPLORE THE FEATURES
Recorded PyTorch example · Click a layer to explore
ResidualClassifier
Input[1, 64]→StemLinear · 64 → 128→ResidualBlock4 operations inside→HeadLinear · 64 → 10→OutputModel result
SelectedModuleRecorded failureChoose a layer to look inside
Browse the model as a list
Follow the shapes.
Inspect the input and output of each operation from a captured run. See where a tensor changes shape, where execution fails, and what was never reached.
Find the exact Python.
Select a layer to inspect its parameters, then choose View source. The graph stays connected to the declaration that defines it.
THE ACTUAL WORKFLOW
From your file to a clearer view.
UNDERSTAND THE CHANGE, TOO
Review architecture changes.
In code and in the graph.
Compare Git revisions or open a GitHub pull request directly in the Canvas. Find the changed layers in the graph, then inspect the Python diff beside them.
Widen the Linear layers from 32 to 64, replace ReLU with GELU, and add LayerNorm. See old and new values on the layers themselves, with added and removed connections showing how the forward pass changes.
~ Modified+ Added− Removed
Open graph diffFor edits made from the graph, review the proposed Python in VS Code. You choose Apply or Discard.
A SHARED UNDERSTANDING
Connect your agent.
Share your model.
Give your coding agent the model’s structure. Give a teammate a graph they can explore. Keep the context attached to the architecture.
MCP
WITH YOUR CODING AGENT
Give your agent model context.
Connect your existing agent through MCP to inspect captured operations and shapes, help with model setup, or propose a change for native review.
Connect your agent Your agent is optional. Core graph tools work without one.WITH YOUR TEAM
Share an interactive model graph.
Export a portable snapshot with saved views and notes. A teammate can explore the architecture and join the discussion in a browser, without Python or VS Code.
An exchanged graph file; your source and weights stay separate.BRING YOUR OWN MODEL
Visualize your
own model.
Open your PyTorch project. Keep your editor, environment, and workflow.
Install in VS Code macOS, Windows & Linux · Remote SSHRequirements & support
Also available on Open VSX
BEFORE YOU START
A few useful details.
What do I need to use tensorViz?
Desktop VS Code 1.106 or newer and Python 3.10–3.14 with PyTorch and your model’s dependencies. Use an existing Conda or virtual environment; the separate tensorViz service needs Python 3.11+. See tested versions and limits ↗
Does opening a graph run my model?
Open as Graph resolves configuration and inputs, then captures one forward call in your selected Python environment. It describes that input and mode. Choose Inspect source only to browse without execution. The website’s examples and shared graph files are snapshots; they never run Python.
Can I work on a remote machine?
Yes. With desktop VS Code Remote SSH, install tensorViz on the SSH host and select Python there. Your files and model execution stay on the server; the Canvas appears on your desktop. WSL, dev containers, and Codespaces are outside the current support boundary.
Do I need an AI subscription or API key?
The graph, input setup, and source-edit forms work without an agent or provider key. You can optionally connect an existing coding agent through MCP. That agent uses its own account and configuration.
What happens when I edit from the graph?
Supported edits become a proposed Python change in a native VS Code diff. Apply changes the unsaved buffer; Discard writes nothing. Save and recapture to check the changed model. Reviewing or applying a change does not prove it runs correctly.
What does a shared graph contain?
A .tensorviz.json snapshot contains the graph,
available capture details, layout, saved views, and notes.
Recipients explore it in the static viewer. It is an exchanged
file, with no live workspace sync; model source and weights are
not bundled.