Batch-wise image augmentation with dynamic parameters in PyTorch

vision
1

I’m currently working on a code for automated data augmentation in PyTorch and I’m wondering if there’s a method to apply a set of augmentations with varying parameters to an entire batch at once.
Here’s how I’m currently implementing it using a loop:

# Assuming we sample a list of parameters using a parameter search algorithm
theta0 = [1, 0, 1, ..., 0]
theta1 = [58, 75, 44, ..., 14]

# Apply augmentation for each image in each batch
...
for i, (inputs, targets) in enumerate(data_loader):
    transformed_imgs = []
    for j in range(len(inputs)):
        transform = transforms.Compose([
            transforms.RandomHorizontalFlip(theta0[j]),
            lambda x: transforms.functional.rotate(x, theta1[j])
        ])
        im = inputs[j]
        transformed_imgs.append(transform(im))

To the best of my knowledge, PyTorch enables the application of augmentations with either entirely fixed or entirely random parameters through the torchvision.transforms module.
I came across this post from a few years ago, where the response suggests applying augmentations per sample.
Is there an efficient way to achieve this kind of functionality?

Thanks in advance!