Skip to content

Commit 97e2cdb

Browse files
authored
[torchlib] Modify aten_unbind to use None for split_sizes
1 parent 07f3e4c commit 97e2cdb

File tree

1 file changed

+2
-3
lines changed
  • onnxscript/function_libs/torch_lib/ops

1 file changed

+2
-3
lines changed

onnxscript/function_libs/torch_lib/ops/core.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8616,12 +8616,11 @@ def aten_type_as(self: TTensor, other: TTensor2) -> TTensor2:
86168616
return op.CastLike(self, other)
86178617

86188618

8619-
@torch_op("aten::unbind.int")
8619+
@torch_op("aten::unbind.int", trace_only=True)
86208620
def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]:
86218621
"""unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]"""
86228622

8623-
split_sizes = op.Constant(value_int=1)
8624-
return op.SplitToSequence(self, split_sizes, axis=dim, keepdims=False)
8623+
return op.SplitToSequence(self, None, axis=dim, keepdims=False)
86258624

86268625

86278626
@torch_op("aten::unflatten.int", trace_only=True)

0 commit comments

Comments
 (0)