How to include Data Augmentation for Roboflow's SegFormer Colab Notebook?

Hi how can I include data augmentation in Roboflow’s Segformer guide.

I’ve tried following this guide with no avail: GPU and batched data augmentation with Kornia and PyTorch-Lightning — PyTorch Lightning 1.8.0dev documentation

Can anyone show code additions/modifications to the SegFormer guide: Google Colab required to make it work?