DataGenerator fix