LibTorch：tensor.index_select()

LibTorch中的tensor.index_select()方法与PyTorch中的用法类似，作用都是在指定的tensor维度dim上按照index值索引向量。先看一下在LibTorch中的声明：

inline Tensor Tensor::index_select(int64_t dim, const Tensor & index)

主要是两个参数，一个是要选择的维度，对于二维tensor来说，0代表按行索引，1代表按列索引，参数index代表索引值，这个参数要重点注意一下！首先index它本身就是一个tensor！另外，还得是int64（kLong）类型的tensor！

下面举例说明：

如果我们要获取一个[shape为(3, 2)]的二维tensor的第列维度上(dim=1)的数据(index=0，即第0列)[shape为(3, 1)]并打印，如果像下面这样写：（from_blob的作用是将数组转化为tensor）

#include "iostream"
#include "torch/script.h"
using namespace torch::indexing;
int main()
{
    float a[3][2] = {1};
    at::Tensor b = at::from_blob(a, {3, 2}, at::kFloat).index_select(1, at::tensor(0));
    std::cout << b << std::endl;
    return 0;
}

如果就这么运行，那么bug就会出现了。。。哈哈哈，先看看bug说的啥吧：

terminate called after throwing an instance of 'c10::Error'
  what():  index_select(): Expected dtype int64 for index (index_select_out_cpu_ at ../aten/src/ATen/native/TensorAdvancedIndexing.cpp:387)

发现问题没？上面不是说了嘛，index它本身就是一个tensor！另外，还得是int64（kLong）类型的tensor！所以我们还要对at::tensor(0)做一下数据类型的转换，如下面这段代码就正常了：

#include "iostream"
#include "torch/script.h"
using namespace torch::indexing;
int main()
{
    float a[3][2] = {{1,2},{3,4},{5,6}};
    at::Tensor b = at::from_blob(a, {3, 2}, at::kFloat).index_select(1, at::tensor(0).toType(at::kLong));
    std::cout << b << std::endl;
    return 0;
}

/*****************运行结果*******************/
 1
 3
 5
[ CPUFloatType{3,1} ]