sklearn & pytorch: Train test split for neural net training in pipeline for a grid search

Question

I am working on a pretty big dataset, which we decided to groupKfold (as we got measurements in the dataset which shouldn\'t get split, but folded i k folds).

We then