DataLoader where each batch belongs to a group

Hello, I have a dataset containing images of different patients. This dataset has an image_id and a patient_id. One patient can have many images. I wish to create sort of a DataLoader where each batch only has images from a single patient. By using a DataLoader in the classic way each batch may contain images from different patients and my Loss function is computed per patient.

One way I thought to solve this is to group by the patient_id, iterate over groups and at each iteration create a new DataLoader but i don’t know if it is the correct way.

Could you tell me a way to load images in the way described?