ValueError in Scenario Data Loading example

[bellma-lilly]

I'm following the Scenario Data Loading example here: https://waymo-research.github.io/waymax/docs/notebooks/data_demo.html

At this line:

# Using logged trajectory
img = visualization.plot_simulator_state(scenario, use_log_traj=True)

I get this error:

ValueError: cannot reshape array of size 1080000 into shape (300,300,3)

How can I solve this issue?

[nononowow]

I got the same error on my m2, any solutions?