Here’s how we can take advantage of the RandAugment class from tf-models.
First, install the nightly build: !pip install -q tf-models-nightly.
Then -
from official.vision.beta.ops import augment

# Recommended is m=2, n=9
augmenter = augment.RandAugment(num_layers=3, magnitude=10)

dataset = load_dataset(filenames)
dataset = dataset.shuffle(batch_size*10)
dataset = dataset.map(augmenter.distort, num_parallel_calls=AUTO)
...

load_dataset() takes image filepaths and reads images from those paths and creates a tf.data.Dataset object.