Access to each training example's original index in the input Anndata object

I am training scvi with augmentations. For one of my augmentations, I need access to each cell’s original index in the full input adata object. What is the best way to go about this?