It would be nice to log the rejections of the background classes during training.
See https://gitlab.cern.ch/atlas-flavor-tagging-tools/algorithms/GNNJetTagger/-/blob/main/gnn_tagger/training/metrics.py