Hi,
It makes more sense in this scenario to write a custom dataloader since nifti is not supported as an extension yet.
Check out PyTorch Custom DataLoader.
I personally use CSV for DataLoading instead.