**Is your feature request related to a problem? Please describe.** There are currently no data augmentation happening during training. **Describe the solution you'd like** Enable augmentations via [TorchIO](https://github.com/fepegar/torchio). **Describe alternatives you've considered** N.A.