Example notebook does not work

[Skylion007]

When trying to run the cifar10 notebook, some of the cells do not seem to work, specifically, finalize_score even on old versions of trak do not work:

/content/journey-TRAK/examples/CIFAR10_trak_results/36/_is_featurized.mmap, the grads.mmap is also incorrectly prefixed. Everything works up until finalize_score in the notebook, unsure what I am missing here.