From 531242be79f638e3fddf4b0bcc9d813f9341a3b4 Mon Sep 17 00:00:00 2001
From: Kaipeng Deng <dengkaipeng@baidu.com>
Date: Thu, 17 Dec 2020 17:49:18 +0800
Subject: [PATCH] fix fields lost (#1914)

---
 ppdet/data/reader.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ppdet/data/reader.py b/ppdet/data/reader.py
index 79133abe9..9be283e51 100644
--- a/ppdet/data/reader.py
+++ b/ppdet/data/reader.py
@@ -105,6 +105,9 @@ class BaseDataLoader(object):
             self._batch_transforms = Compose(batch_transforms,
                                              copy.deepcopy(self._fields),
                                              transform, num_classes)
+            self.output_fields = self._batch_transforms.output_fields
+        else:
+            self.output_fields = self._fields
 
         self.batch_size = batch_size
         self.shuffle = shuffle
@@ -161,7 +164,7 @@ class BaseDataLoader(object):
         # data structure in paddle.io.DataLoader
         try:
             data = next(self.loader)
-            return {k: v for k, v in zip(self._fields, data)}
+            return {k: v for k, v in zip(self.output_fields, data)}
         except StopIteration:
             self.loader = iter(self.dataloader)
             six.reraise(*sys.exc_info())
-- 
GitLab