fix: op.function doesn't work with slices.

deepmodeling · May 25, 2024 · e5d56fe · e5d56fe
1 parent 8c8c1b7
commit e5d56fe
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/dflow/python/python_op_template.py b/src/dflow/python/python_op_template.py
@@ -500,8 +500,9 @@ def render_script(self):
         script += "op_obj.workflow_name = '{{workflow.name}}'\n"
         script += "if __name__ == '__main__':\n"
         script += "    input = OPIO()\n"
-        script += "    input_sign = %s.get_input_sign()\n" % class_name
-        script += "    output_sign = %s.get_output_sign()\n" % class_name
+        script += "    from copy import deepcopy\n"
+        script += "    input_sign = deepcopy(%s.get_input_sign())\n" % class_name
+        script += "    output_sign = deepcopy(%s.get_output_sign())\n" % class_name
         if self.slices is not None and self.slices.pool_size is not None:
             script += "    from dflow.python.utils import try_to_execute\n"
             script += "    from functools import partial\n"
@@ -572,7 +573,6 @@ def render_script(self):
                     % name
             script += "    assert n_slices is not None\n"
             script += "    input_list = []\n"
-            script += "    from copy import deepcopy\n"
             script += "    for i in range(n_slices):\n"
             script += "        input1 = deepcopy(input)\n"
             for name in sliced_inputs: