KeyError: 'CollectiveSize' when running Analysis Jupyter script

[erwincoumans]

I'm trying to generate a flamegraph using the generated install.sh script from nsys-jax. The install.sh script in the zip file opens a Jupyter notebook, and when running the Analysis.ipbn script, some error shows up:

KeyError: 'CollectiveSize', any idea?
---> 34 max_collective_size = comm_df["CollectiveSize"].max()
     35 assert (
     36     num_profiled_devices == max_collective_size

Any idea? Thank you!

Host: Ubuntu 22.04, Python 3.11.7.
Attached also the zip file, in case it helps.
nsys-jax-report-c9dgz7dt.zip

[olupton]

Should be fixed by #988, which will be available in the 2024-08-10 containers and later.