Have a look at @christianperone’s …MedicalTorch repo.
It might be a good starter for your project.
Regarding the data augmentations, you could try to apply the augmentation on each slide of your scans.
In this case I would use the functional API of torchvision.transforms to make sure that each “random” transformation is applied in with the same parameters on each slide.
Here is a small example on using the same transform parameters on the data and target. You would need to adapt this code for all image slides.
Alternatively you could write your own (affine) transformations. Here is a small example I’ve written for a rotation using rotation matrices.