helios: Fix to train test partition in case of long X Ray wavelengths
helios: Fix to train test partition in case of long X Ray wavelengths

transcendent-ai-labs/PlasmaML,transcendent-ai-labs/PlasmaML,mandar2812/PlasmaML,transcendent-ai-labs/PlasmaML,mandar2812/PlasmaML