Ok, I would suggest to overload the Dataset class for your own implementation of loading and preprocessing the data.
Since you are creating your own Dataloader, you are losing the functionality of the built-in DataLoader 54, which uses multiprocessing, shuffling etc.
I couldn’t follow your code completely, but have a look at the Data Loading and Processing Tutorial 262, which gives you all necessary information how to re-implement your code.