Preprocessing & augmentation during training

Hi! I think Dataset.set_transform is what you are looking for.