Augmentation pipeline

It would be nice to be able to apply transformations on data on the fly as it is loaded. This could be used to add systematics to the model.

Some things to consider

• the augmentations should happen during training only
• it might be cleaner to add them in the dataloading stage rather than inside the model in an if (up for debate)
• we eventually would want to support different augmentations e.g. using gaussians or parameterised maps