Solved by ptrblck in post #2 
              
              
                Usually you could create a separate Dataset instance for the training, validation, and test splits. 
Based on your code snippet you could pass an argument to the __init__ method and use it to create the corresponding split (e.g. by passing train, dev, or test).