Fix op index_select (#756)

Summary:
Pull Request resolved: #756

ghstack-source-id: 203458981
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D50104697

fbshipit-source-id: 4209839556311ea20141d8dcaa585bc70d47d9a3

Author: manuelcandales

kernels/portable/cpu/op_index_select.cpp

@@ -28,10 +28,15 @@ Tensor& index_select_out(
   ET_KERNEL_CHECK(
       ctx, check_index_select_args(in, dim, index, out), InvalidArgument, out);
 
+  if (dim < 0) {
+    dim += nonzero_dim(in);
+  }
+
   size_t expected_ndim = 0;
   Tensor::SizesType expected_size[kTensorDimensionLimit];
   get_index_select_out_target_size(
       in, dim, index, expected_size, &expected_ndim);
+
   ET_KERNEL_CHECK(
       ctx,
       resize_tensor(out, {expected_size, expected_ndim}) == Error::Ok,

kernels/portable/cpu/util/index_util.cpp

@@ -279,7 +279,7 @@ bool check_index_select_args(
 
   if (index.scalar_type() == ScalarType::Long) {
     const int64_t* const index_ptr = index.const_data_ptr<int64_t>();
-    for (size_t i = 1; i < index.numel(); ++i) {
+    for (size_t i = 0; i < index.numel(); ++i) {
       ET_LOG_MSG_AND_RETURN_IF_FALSE(
           index_ptr[i] >= 0 && index_ptr[i] < nonempty_size(in, dim),
           "index[%zu] = %" PRId64 " is out of range [0, %zd)",
@@ -289,7 +289,7 @@ bool check_index_select_args(
     }
   } else {
     const int32_t* const index_ptr = index.const_data_ptr<int32_t>();
-    for (size_t i = 1; i < index.numel(); ++i) {
+    for (size_t i = 0; i < index.numel(); ++i) {
       ET_LOG_MSG_AND_RETURN_IF_FALSE(
           index_ptr[i] >= 0 && index_ptr[i] < nonempty_size(in, dim),
           "index[%zu] = %" PRId32 " is out of range [0, %zd)",