fix: pass on sliced big parameter for sliced super OP

Signed-off-by: zjgemi <[email protected]>
deepmodeling · May 29, 2024 · a2e1867 · a2e1867
1 parent b4f5452
commit a2e1867
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/dflow/step.py b/src/dflow/step.py
@@ -2288,8 +2288,8 @@ def add_slices(templ: OPTemplate, slices: Slices, layer=0):
         for step in steps:
             for par in list(step.inputs.parameters.values()):
                 # input parameter referring to sliced input parameter
-                if getattr(par, "value", None) is \
-                        templ.inputs.parameters[name]:
+                if getattr(par, "value", None) is templ.inputs.parameters[
+                        name] or par.source is templ.inputs.parameters[name]:
                     templ.inputs.parameters["dflow_skip_slice_input"] = \
                         InputParameter(value=0)
                     if hasattr(step.template, "set_skip_slice_input"):