Update onnxscript/function_libs/torch_lib/ops/core.py

Co-authored-by: Copilot <[email protected]>

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -8767,7 +8767,7 @@ def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]:
     if isinstance(self.shape[dim], int):
         # We can create a definitive split op if the input shape is static
         outputs = op.Split(self, axis=dim, num_outputs=self.shape[dim])
-        return [op.Squeeze(out, [self.shape[dim]]) for out in outputs]
+        return [op.Squeeze(out, [dim]) for out in outputs]
 
     return op.SplitToSequence(self, axis=dim, keepdims=False)