Your approach of using a custom Dataset sounds correct.
To implement it, you could have a look at this tutorial 4 first and adapt it so that it fits your use case.
In particular, you would store the image/patient input images and the corresponding mask paths in the __init__ method and would then lazily load each pair in the __getitem__.
The critical step would be to make sure the paths for the DICOM images and masks are sorted correctly, so that indexing these lists of paths yields the desired corresponding pairs and doesn’t mix patients etc.
Also, you could take a look at e.g. MONAI 2, which might already provide some Dataset implementations for similar use cases.