No text label displayed on Google Colab (Pytorch 1.8.1)

[makaishi2]

I run following sample code on both local anaconda and google colab.
And on google colab no text label dispkayed on a result node.

Sample Code:

import torch
from torchviz import make_dot
print(torch.__version__)
x_np = np.arange(-2, 2.1, 0.25)
x = torch.tensor(x_np, requires_grad=True)
y = 2 * x**2 + 2
z = y.sum()
params = dict(x = x)
make_dot(z, params=params)

Result 1 (local anaconda pytorch 1.7.1)

Result 2 (Google Colab pytorch 1.8.1)

[albanD]

Hi,

As mentioned in the doc here:

pytorchviz/torchviz/dot.py

Line 50 in 5c91c3d

- Green: tensor passed in as outputs

The green node is the Tensor passed as output. The () is its size.
You can pass z=z to the params dict and that name will be displayed there.