The notebook neuralcollapse/neuralcollapse.ipynb contains code demonstrating Neural Collapse on Cross-Entropy and MSE Loss. Notebook is designed to be short, easy-to-interpret, and executable from the browser using Google Colab. (There is a link in the notebook to open in Colab.)

MNIST-ResNet18 was chosen because it ran most reliably within the in-browser the memory constraints of Google Colab. If you are still getting out-of-memory errors, try clicking "Runtime"->"Factory Reset Runtime" on the Colab menu bar.

It should be clear how to adapt the code to other networks-dataset combinations to be run on local clusters with more memory. Happy experimenting!