Add option to retrieve fold info

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