Solved by SreenivasVRao in post #6 
              
              
                Yes exactly! At the end of the process, I get one batch, and the sampled value k. 
I achieved this by delaying I/O and preprocessing until the collate function, but it’s not very flexible. 
at the moment, I do: 
def collate_data(batch):
    k = np.random.randint(0, 10)
    batch = [sample[k] for sam…