fixed a typo

1 files changed, 1 insertions, 1 deletions

diff --git a/src/graph/expression_operators.cpp b/src/graph/expression_operators.cpp
index db820062..826bd9f0 100755
--- a/src/graph/expression_operators.cpp
+++ b/src/graph/expression_operators.cpp
@@ -256,7 +256,7 @@ Expr index_select(Expr a, int axis, Expr indices) {
   // We have specialized kernels for non-batched indexing of first or last axis of a 2D tensor.
   auto rank = a->shape().size();
   if (rank == 2) {
-    if (axis == 0 || axis == 2)
+    if (axis == 0 || axis == -2)
       return Expression<RowsNodeOp>(a, indices);
     else if (axis == -1 || axis == 1)
       return Expression<ColsNodeOp>(a, indices);