Multiple indexing in C++ eg: data[0,1,2]

I downloaded the c++ preview build of pytorch and am trying to do multiple indexing via the [] operator:

in python it would be:

data = torch.rand(3,3,3)
data[tensor([0,1]), tensor([1,2])]

Anyone know the syntax for this functionality in C++? in C++ the [] operator only supports a single int.

Found it for anyone wondering. Just call index:

Tensor data = at::rand(3,3,3);
dim1 = at::tensor(({0, 1}, torch::kInt64);
dim2 = at::tensor(({1, 2}, torch::kInt64);
data.index({dim1, dim2});
1 Like