Problems on nii dataload

I am recently working on organ segmentation that needs nii files process. It is fine converting nii to img then load them to train, but as everyone knows medical imaging process doesn’t regard it an optimal choice, I have to directly handle nii to avoid information loss.
Problem’s here. Nii files to tensors will stress my device out even though I am working on COLAB. So I am now dying to know how can I better load nii files with fastai? Thanks soooooooooo much if anyone could help me.

class NiiDataset(Datasets):
    def __init__(self, volume_list, segmentation_list, window_type="liver", slice_interval=2):
        self.volume_list = volume_list
        self.segmentation_list = segmentation_list
        self.window_type = window_type
        self.slice_interval = slice_interval

        # Load NIfTI files using read_nii and select slices with interval
        self.volumes = [read_nii(path)[::self.slice_interval] for path in volume_list]
        self.masks = [read_nii(path)[::self.slice_interval] for path in segmentation_list]

    def __len__(self):
        return len(self.volumes)

    def __getitem__(self, idx):
        volume = torch.from_numpy(self.volumes[idx])
        mask = torch.from_numpy(self.masks[idx])

        # Apply windowing to volume slices only
        w, l = dicom_windows.__dict__[self.window_type]
        volume = torch.stack([windowed(slice, w, l) for slice in volume])

        return volume.float(), mask.long()  # Specify data types
codes = np.array(["background", "liver", "tumor"])

data_block = DataBlock(blocks=(NiiDataset(volume_list, segmentation_list), NiiDataset(volume_list, segmentation_list)), 
                       get_x=lambda x: x[0], 
                       get_y=lambda x: x[1],
                       codes=codes)  # Provide class codes

dls = data_block.dataloaders(volume_list, segmentation_list, bs=4)

I have never done any medical imaging work so I can’t speak from experience but I did come across this link so I’m dropping it here in case you haven’t seen it: