Visualization in tutorial not working in colab and remote setup

I have tried running this tutorial both in google colab and in jupyterlab on a remote server setup.

https://ax.dev/tutorials/tune_cnn.html

Both times the render()-call returns an empty output.

This is the code, no changes from tutorial except installation of ax:

!pip install ax-platform
%%
import torch
import numpy as np

from ax.plot.contour import plot_contour
from ax.plot.trace import optimization_trace_single_method
from ax.service.managed_loop import optimize
from ax.utils.notebook.plotting import render, init_notebook_plotting
from ax.utils.tutorials.cnn_utils import load_mnist, train, evaluate, CNN

init_notebook_plotting()
%%
torch.manual_seed(12345)
dtype = torch.float
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
%%
BATCH_SIZE = 512
train_loader, valid_loader, test_loader = load_mnist(batch_size=BATCH_SIZE)
%%
def train_evaluate(parameterization):
    net = CNN()
    net = train(net=net, train_loader=train_loader, parameters=parameterization, dtype=dtype, device=device)
    return evaluate(
        net=net,
        data_loader=valid_loader,
        dtype=dtype,
        device=device,
    )
%%
best_parameters, values, experiment, model = optimize(
    parameters=[
        {"name": "lr", "type": "range", "bounds": [1e-6, 0.4], "log_scale": True},
        {"name": "momentum", "type": "range", "bounds": [0.0, 1.0]},
    ],
    evaluation_function=train_evaluate,
    objective_name='accuracy',
)
%%
render(plot_contour(model=model, param_x='lr', param_y='momentum', metric_name='accuracy'))

switching to classic jupyter notebook (in remote setup) as suggested in https://github.com/facebook/Ax/issues/94 doesn’t change anything.

pip freezes for both environments:

colab_pip_freeze.txt

jupyterlab_pip_freeze.txt

[saitcakmak]

Here’s a simple workaround. Add

import plotly.io as pio
pio.renderers.default = "colab"

to the NB running on colab. The figure renders and is interactive 😃 (for smaller notebook file size at the expense of non-interactive plots, you can use “png” as the renderer)

Note: You also need to do !pip install kaleido.

[2timesjay]

Still investigating colab half of this issue (realized this one was the right one to host discussion on).