Skip to content

Commit 4042df3

Browse files
authored
Add missing output_size kwarg to repeat_interleave (#2691)
Fixes #2687 --------- Signed-off-by: Yuan Yao <yuanyao@nvidia.com>
1 parent 1dd9d04 commit 4042df3

File tree

1 file changed

+4
-1
lines changed
  • onnxscript/function_libs/torch_lib/ops

1 file changed

+4
-1
lines changed

onnxscript/function_libs/torch_lib/ops/core.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7539,7 +7539,10 @@ def aten_repeat(self: TTensor, repeats: Sequence[TInt]) -> TTensor:
75397539

75407540
@torch_op("aten::repeat_interleave.self_int", trace_only=True)
75417541
def aten_repeat_interleave_self_int(
7542-
self: TensorType, repeats: int, dim: Optional[int] = None
7542+
self: TensorType,
7543+
repeats: int,
7544+
dim: Optional[int] = None,
7545+
output_size: Optional[int] = None,
75437546
) -> TensorType:
75447547
"""repeat_interleave.self_int(Tensor self, SymInt repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor
75457548

0 commit comments

Comments
 (0)