Adjust AutoAugment transforms?

I’m trying to apply some random augmentations to my images and found there is transforms.AutoAugment
which can help.

However, my data needs to eliminate rotations from the augment, or at least make it so only very minor rotations can be applied.

Can the AutoAugment somehow be modified to do what I require?

You could do something like subclass the existing the existing AutoAugment class and add an if statement to the forward method to skip the rotation augmentation based on the op_name in the for loop.

https://pytorch.org/vision/main/_modules/torchvision/transforms/autoaugment.html#AutoAugmentPolicy