How do you do data augmentation for segmentation in PyTorch

Question

I am trying to add augmentation to my data loader in torch. All of my example images have masks and none are blank. But when I add an augmentation like RandomRotation<