How to implement test time augmentation with pytorch lightning

I’ve read about test time augmentation but how would I implement this using say torchvision transforms/albumentations and the LightningModule?