Merge pull request #1094 from onnx/tom/FixSatefulPartitionedCallError

Added f to the list of TF function attributes
onnx · Sep 8, 2020 · 0a11159 · 0a11159
2 parents 19d3f97 + d8bc10b
commit 0a11159
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tf2onnx/tf_utils.py b/tf2onnx/tf_utils.py
@@ -203,7 +203,7 @@ def tflist_to_onnx(g, shape_override, const_node_values=None):
             elif a == "output_shapes":
                 # we should not need it since we pull the shapes above already
                 pass
-            elif a in {"body", "cond", "then_branch", "else_branch"}:
+            elif a in {"body", "cond", "then_branch", "else_branch", "f"}:
                 input_shapes = [inp.get_shape() for inp in node.inputs]
                 nattr = get_tf_node_attr(node, a)
                 attr[a] = nattr.name