Hey @ptrblck,
I think I have understood the issue. Since I fed only patches from 2 images, thus the batch never changes beyond 2. Sorry for the confusion.
Now the questions remain,
I am giving Dataset one image and one segmentation path in index, which loads one image[170,256,256] and segmentation[same size] file from directory. As you suggested, I am generating patches inside __getitem__
after patch operation,
img: [192, 16, 64, 64]
seg: [192, 8, 32, 32]
finally I create segmentation channels for 9 labels
img: [192, 1 , 16, 64, 64]
seg: [192, 9, 8, 32, 32]
This way I can directly load data from the directories and need not save anything.
Is there a way to iterate 192 patches as
src: [batch_size, 1, 16, 64, 64]
trgt:[batch_size, 9, 8, 32, 32] and feed to Dataloader and also iterate 3D images in Dataset?
Now to do the same I had to save all the patches to an hdf5(around 8gb) file as and index one by one.
    g = h5py.File(hpath[0], 'r')
    print(g.keys())
    img = np.array(g['MR'])    # (114, 192, 16, 64, 64)
    seg = np.array(g['Mask'])   # (114, 192, 8, 32, 32)


    train_img = img.reshape(-1, 16, 64, 64)
    train_seg = seg.reshape(-1, 8, 32, 32)

    trainSet = SimpleDataset(train_img, train_seg)
    x, y = next(iter(trainSet))
    print(x.shape, '\n', y.shape)
    # (1, 1, 16, 64, 64)
    # (1, 9, 8, 32, 32)

Is there a more efficient way to do that?
I hope I make sense and convey it clearly