diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index e5157f90ee..2af4b8a4fe 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -41,6 +41,13 @@ def get_shape_from_type_proto(type_proto): return None # note no shape is different from shape without dim (scalar) +def get_elem_type_from_type_proto(type_proto): + if is_sequence(type_proto): + return type_proto.sequence_type.elem_type.tensor_type.elem_type + else: + return type_proto.tensor_type.elem_type + + def get_shape_from_value_info(vi): cls_type = vi.type.WhichOneof("value") if cls_type is None: @@ -570,7 +577,7 @@ class SymbolicShapeInference: vi.CopyFrom( helper.make_tensor_value_info( node.output[0], - self.known_vi_[node.input[0]].type.tensor_type.elem_type, + get_elem_type_from_type_proto(self.known_vi_[node.input[0]].type), self._get_shape(node, 0), ) )