chore: trying to fix cp's shape

zama-ai · Jun 5, 2024 · 893e956 · 893e956
1 parent 746e5c2
commit 893e956
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/concrete/ml/common/preprocessors.py b/src/concrete/ml/common/preprocessors.py
@@ -1571,12 +1571,12 @@ def add_extra_shape_to_subgraph(tlu_node, extra_dim: Tuple[int, ...]):
                 node.evaluator.properties["constant"],
                 node.evaluator.properties["constant"].shape + tuple([1 for _ in extra_dim]),
             )
+            node.output.shape = node.evaluator.properties["constant"].shape
             continue
 
         node.inputs[0].shape = deepcopy(node.inputs[0].shape)
         node.inputs[0].shape = node.inputs[0].shape + extra_dim
 
-        # print(f"{node.output.shape=}")
         node.output.shape = deepcopy(node.output.shape)
 
         if node.output.shape[-1:] == extra_dim:  # BIG HACK