Are we loading the whole dataset in the memory?

Hi! I was trying to adapt the notebook of the downstream task of lung nodule segmentation (link) to brain tumor segmentation (BraTS). I noticed that the whole dataset is loaded in the memory and no pytorch dataloaders are used, nor custom dataloader classes are implemented. Is this the case? Because I am having trouble doing this for BraTS which is a large dataset. Could you perhaps provide the downstream task notebook for BraTS too, so that I can see how you managed to do this?

Start with keras/downstream_tasks/lung nodule segmentation.ipynb, the notebook linked in the issue, and inspect how its dataset is loaded. Check whether the full dataset is held in memory and how PyTorch dataloaders could apply to the BraTS adaptation. Done would require a documented loading approach or a BraTS downstream-task notebook that addresses large-dataset handling.