cifar_quickstart.ipynb is a complete notebook that trains models from scratch and computes TRAK scores using the trained checkpoints. See tutorial for a walk-through.

imagenet.py computes TRAK scores using a pre-trained PyTorch ImageNet model.

qnli.py computes TRAK scores for a BERT model. To run it, you need to supply a checkpoint finetuned on GLUE QNLI. You can get them by running the run_glue.py script from Hugging Face.

See tutorial for a walk-through.