diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 9de7b170f..6fb230e90 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8085,7 +8085,7 @@ def aten_swapdims(self: TensorType, dim0: int, dim1: int) -> TensorType: @torch_op("aten::sym_size.int", trace_only=True) def aten_sym_size(self: TensorType, dim: int = 0) -> INT64: """sym_size.int(Tensor self, int dim) -> SymInt""" - return op.Shape(self, end=dim + 1, start=dim) + return op.Squeeze(op.Shape(self, end=dim + 1, start=dim)) def aten_symeig(