fix wrong shapes in loop body inputs if shape invariances are set in TF

Signed-off-by: Salvetti, Francesco <[email protected]>

Author: f-salvetti

tf2onnx/onnx_opset/controlflow.py

@@ -571,7 +571,8 @@ def wire_while_body(parent_g, g, loop_node, body_input_to_state_var, cond_input_
     g.set_dtype(func_inputs[0], onnx_pb.TensorProto.INT64)
     g.inputs = [g.get_node_by_output(inp) for inp in func_inputs]
 
-    for p, c in zip(loop_node.input, func_inputs):
+    # we should use outputs shape, not inputs, since there may be shape invariants
+    for p, c in zip(loop_node.output, func_inputs[2:]):
         g.copy_shape(p, c)
 
     for i, node in enumerate(g.inputs):