2749 Fix ThreadDataLoader reset (#2750)

* [DLMED] fix continuously run

Signed-off-by: Nic Ma <[email protected]>

* [MONAI] python code formatting

Signed-off-by: monai-bot <[email protected]>

* [DLMED] remove out-date comment

Signed-off-by: Nic Ma <[email protected]>

Co-authored-by: monai-bot <[email protected]>

Author: Nic-Ma

monai/data/thread_buffer.py

@@ -87,8 +87,6 @@ class ThreadDataLoader(DataLoader):
     def __init__(self, dataset: Dataset, num_workers: int = 0, **kwargs):
         super().__init__(dataset, num_workers, **kwargs)
 
-        # ThreadBuffer will use the inherited __iter__ instead of the one defined below
-        self.buffer = ThreadBuffer(super().__iter__())
-
     def __iter__(self):
-        yield from self.buffer
+        buffer = ThreadBuffer(super().__iter__())
+        yield from buffer

tests/test_thread_buffer.py

@@ -48,6 +48,8 @@ def test_dataloader(self):
         for d in dataloader:
             self.assertEqual(d["image"][0], "spleen_19.nii.gz")
             self.assertEqual(d["image"][1], "spleen_31.nii.gz")
+
+        for d in dataloader:
             self.assertEqual(d["label"][0], "spleen_label_19.nii.gz")
             self.assertEqual(d["label"][1], "spleen_label_31.nii.gz")