To apply the jitter over the entire dataset, use the 🤗 Datasets [~datasets.Dataset.set_transform] function. The transform is applied on the fly which is faster and consumes less disk space:

train_ds.set_transform(train_transforms)
test_ds.set_transform(val_transforms)