multiprocess exception fix

Same fix as that of pytorch#1228
SarthakYadav · Dec 18, 2024 · cbe5475 · cbe5475
1 parent 792b5e8
commit cbe5475
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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