equivariance_measure/data.py
per_class_augmentation/data.py