From 38f38a9139f96e3514b77b8ccf2ac67dde607b0f Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Tue, 13 Jun 2023 16:04:20 +0800
Subject: [PATCH] fix KeyError when delte idx from task_infos (#54498)

---
 python/paddle/io/dataloader/dataloader_iter.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/paddle/io/dataloader/dataloader_iter.py b/python/paddle/io/dataloader/dataloader_iter.py
index b212dd42f1d..0ffe7c46e77 100644
--- a/python/paddle/io/dataloader/dataloader_iter.py
+++ b/python/paddle/io/dataloader/dataloader_iter.py
@@ -752,7 +752,8 @@ class _DataLoaderIterMultiProcess(_DataLoaderIterBase):
                     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:
-- 
GitLab