An interactive tool to visualize the forward pass of a PyTorch model directly in the notebook—with a single line of code. Works with web-based notebooks like Jupyter, Google Colab and Kaggle. Also allows you to export the visualization as image, svg and HTML.

(e.g., shape mismatches) for ease of debugging

• Step-by-step tutorial 👉 here
• Quick Google Colab tutorial 👉 here (must be logged in to Colab)
• Check out demos 👉 here

Install via pip

pip install torchvista

Alternatively, install via conda (See torchvista-feedstock for more information)

conda install -c conda-forge torchvista

Run from your web-based notebook (Jupyter, Colab, VSCode notebook, etc)

import torch
import torch.nn as nn

# Import torchvista
from torchvista import trace_model

# Define your module
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# Instantiate the module and tensor input
model = LinearModel()
inputs = torch.randn(2, 10)

# Trace!
trace_model(model, inputs)

trace_model(
    model,
    inputs,
    show_non_gradient_nodes=True,
    collapse_modules_after_depth=1,
    forced_module_tracing_depth=None,
    height=800,
    width=None,
    export_format=None,
    show_module_attr_names=False,
    export_path=None,
    show_compressed_view=False,
)