Simple addition to the HepNet predict() function to retrieve information about which event was train/val/test in which fold.