Solved by jiwidi in post #3 
              
              
                Thanks for the reply! 
Really good material you linked to, I think I have solved it. Can you double-check my logic? I tested it and works good so far. 
I replace the dataset with this new definition: 
class CustomIterableDatasetv1(IterableDataset):

    def __init__(self, filename):

        #Store …