!pip install muaraugment

from muar.augmentations import BatchRandAugment, MuAugment

# muar augmentations
rand_augment = BatchRandAugment(N_TFMS=3, MAGN=4)
mu_augment = MuAugment(rand_augment, N_COMPS=4, SELECTED=2)

# model
model = LitClassifier(mu_augment)

# data
train, val, test = mnist()

# train
trainer = Trainer()
trainer.fit(model, train, val)