slice matrix and higher order tensors

[YifanShenSZ]

tensor.index works fine with vectors, but thing gets weird for matrix and higher order tensors. An example is:

#include
#include <torch/torch.h>

int main() {
at::Tensor H = at::tensor({1.0, 2.0, 3.0,
4.0, 5.0, 6.0,
7.0, 8.0, 9.0});
H.resize_({3, 3});

at::Tensor index_tensor = at::empty(2, c10::TensorOptions().dtype(torch::kInt64));
index_tensor[0] = 0;
index_tensor[1] = 2;
at::indexing::TensorIndex index(index_tensor);
std::cout << H.index(index) << '\n';

std::vector<at::indexing::TensorIndex> indices_vec(2, index);
c10::ArrayRef<at::indexing::TensorIndex> indices(indices_vec);
std::cout << H.index(indices) << '\n';

at::Tensor H_new = H.clone();
std::cout << H_new.index({0, 2}).index({0, 2}) << '\n';

}

Whose output is:

1 2 3
7 8 9
[ CPUDoubleType{2,3} ]
1
9
[ CPUDoubleType{2} ]
terminate called after throwing an instance of 'c10::IndexError'
what(): too many indices for tensor of dimension 0

i.e. c10::ArrayRefat::indexing::TensorIndex gives {H[0][0], H[2][2]} instead of desirable H[[0,2], [0,2]], while call torch.index twice always acts on dimension 0