fix: set output parameter slices

Signed-off-by: zjgemi <[email protected]>
deepmodeling · Dec 5, 2024 · 5b388ab · 5b388ab
1 parent 43cb6d7
commit 5b388ab
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/dflow/python/python_op_template.py b/src/dflow/python/python_op_template.py
@@ -601,6 +601,9 @@ def render_script(self):
             if isinstance(sign, Artifact):
                 slices = self.get_slices(output_artifact_slices, name)
                 script += "    op_obj.slices['%s'] = %s\n" % (name, slices)
+            else:
+                slices = self.get_slices(output_parameter_slices, name)
+                script += "    op_obj.slices['%s'] = %s\n" % (name, slices)
 
         script += "    import signal\n"
         script += "    def sigterm_handler(signum, frame):\n"