TrivialAugmentWide on segmentation

Hi, what is correct process to apply TrivialAugmentWide to both on input and label when I want to perform segmentation task ?

I guess a safe approach would be creating a custom transform module accepting two inputs (the data input and segmentation mask) by either deriving from the transformation directly or by copying the source code. This would allow you to keep the randomness logic in the forward and then to apply the actual transformation op on both tensors.
Re-seeding the code could also work but could be tricky to get right as you would need to guarantee the same calls into the pseudorandom number generator and could have side effects on the actual training.