From cbe54759ba4c2044a2e843c3d7b4ef2faa8e9b1f Mon Sep 17 00:00:00 2001 From: Sarthak Yadav Date: Wed, 18 Dec 2024 21:10:30 +0100 Subject: [PATCH] multiprocess exception fix Same fix as that of https://github.com/pytorch/data/issues/1228 --- torchdata/dataloader2/communication/iter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchdata/dataloader2/communication/iter.py b/torchdata/dataloader2/communication/iter.py index 73a54d1c9..dcbeca191 100644 --- a/torchdata/dataloader2/communication/iter.py +++ b/torchdata/dataloader2/communication/iter.py @@ -431,17 +431,17 @@ def __iter__(self): disabled_pipe[res_idx] = True cnt_disabled_pipes += 1 disabled = True - req_idx = next(req_idx_cycle) else: # Only request if buffer is empty and has not reached the limit if len(self.res_buffers[res_idx]) == 0 and ( self._limit is None or self._request_cnt < self._limit ): self.datapipes[req_idx].protocol.request_next() - req_idx = next(req_idx_cycle) + self._request_cnt += 1 total_req_cnt += 1 total_res_cnt += 1 + req_idx = next(req_idx_cycle) res_idx = next(res_idx_cycle) if not disabled: yield response.value