diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 869a048511..2c50b4a6b0 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -1331,6 +1331,7 @@ class SymbolicShapeInference: vi = self.known_vi_[node.output[0]] elem_type = self.known_vi_[node.input[0]].type.tensor_type.elem_type vi.type.tensor_type.elem_type = elem_type + vi.type.tensor_type.shape.CopyFrom(onnx.TensorShapeProto()) if len(node.output) > 1: data_shape = self._get_shape(node, 0) diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 6ba3a091ac..dd4ed8b437 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -134,6 +134,32 @@ class TestSymbolicShapeInferenceForOperators(unittest.TestCase): ] self._check_shapes(graph, inferred.graph, expected_shapes) + def test_softmax_cross_entropy_loss(self): + hidden_size = 1024 + + nodes = [ + helper.make_node("SoftmaxCrossEntropyLoss", + inputs=["logits", "labels"], + outputs=["loss"]), + ] + + inputs = [ + helper.make_tensor_value_info('logits', TensorProto.FLOAT, ['b', 's', hidden_size]), + helper.make_tensor_value_info('labels', TensorProto.INT32, ['b', 's']), + ] + + outputs = [ + helper.make_tensor_value_info('loss', TensorProto.FLOAT, None), + ] + + graph = helper.make_graph(nodes, "SoftmaxCrossEntropyLoss_Test", inputs, outputs, []) + model = helper.make_model(graph) + + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + expected_shapes = [ + helper.make_tensor_value_info('loss', TensorProto.FLOAT, []) + ] + self._check_shapes(graph, inferred.graph, expected_shapes) class TestSymbolicShapeInferenceForSlice(unittest.TestCase): def check_slice_of_concat(self, input_dims, start, end, step, expected_output_dim):