I don’t know exactly how this code snippet is supposed to work:
        # Iterate over data
        loop = tqdm(batch, total=len(batch), leave=False)
        images, patient = batch[0]['image'], batch[0]['patient']

but it seems as if you are using a single sample only?
Could you explain this approach a bit more and how you are iterating the entire dataset?