fix KeyError when delte idx from task_infos (#54498)

PaddlePaddle · Jun 13, 2023 · 38f38a9 · 38f38a9
1 parent b5fe3f6
commit 38f38a9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/paddle/io/dataloader/dataloader_iter.py b/python/paddle/io/dataloader/dataloader_iter.py
@@ -752,7 +752,8 @@ def _get_data(self):
                     batch.reraise()
 
                 if idx == self._rcvd_idx:
-                    del self._task_infos[idx]
+                    if idx in self._task_infos:
+                        del self._task_infos[idx]
                     self._structure_infos.append(structure)
                     return batch
                 else: