Fix impute when output_mapping is not empty (#1015)

argilla-io · Oct 7, 2024 · 4b8903b · 4b8903b
1 parent d5c0484
commit 4b8903b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/distilabel/steps/base.py b/src/distilabel/steps/base.py
@@ -591,7 +591,7 @@ def impute_step_outputs(
         result = []
         for row in step_output:
             data = row.copy()
-            for output in self.outputs:
+            for output in self.get_outputs().keys():
                 data[output] = None
             result.append(data)
         return result

diff --git a/src/distilabel/steps/tasks/base.py b/src/distilabel/steps/tasks/base.py
@@ -127,7 +127,7 @@ def impute_step_outputs(
         result = []
         for row in step_output:
             data = row.copy()
-            for output in self.outputs:
+            for output in self.get_outputs().keys():
                 data[output] = None
             data = self._maybe_add_raw_input_output(
                 data,