Reorder some logic

Author: jignparm

tf2onnx/onnx_opset/tensor.py

@@ -1064,16 +1064,19 @@ def version_1(cls, ctx, node, **kwargs):
             axis += len(shape)
         # split the tensor into n outputs
         node.type = "Split"
-        output_shape = ctx.get_shape(node.output[0])
-        if output_shape:
-            ctx.set_shape(node.output[0], output_shape.insert(axis, 1))
+
         # for each output we need to squeeze axis
         for n in node.output:
             op_name = utils.make_name(node.name)
             squeeze_node = ctx.insert_new_node_on_output("Squeeze", n, name=op_name, axes=[axis])
             ctx.copy_shape(n, squeeze_node.output[0])
             ctx.copy_dtype(n, squeeze_node.output[0])
 
+        # split node is 1 rank higher than squeeze nodes
+        output_shape = ctx.get_shape(node.output[0])
+        if output_shape:
+            ctx.set_shape(node.output[0], output_shape.insert(axis, 1))
+
 
 @tf_op("OneHot")
 class OneHot: