Comparable C++ Function to the PyTorch torch.index_select()

C++
1

I have a tensor of shape 300 x 500 x 40 x 40 and a list of indices (e.g. length of list is 100 with values between 0 and 499)

With tensor = torch.index_select(tensor, 1, torch.tensor(index_list)) I reduce my tensor to a shape of 300 x 100 x 40 x 40 by only selecting the values for the indices at dimension 1 (500 → 100).

What is the comparable function in the C++ API to achieve the same?

So far my workaround is the following:

torch::Tensor features;
torch::Tensor features_reduced;
for (int i = 0; i < indices_list.sizes()[0]; ++i) {
  if (i == 0) { // get first
    features_reduced = features.index({torch::indexing::Slice(), indices_list[i] , torch::indexing::Slice(), torch::indexing::Slice()}).unsqueeze(1);
  }
  else { // concatenate to the previous tensor
    features_reduced = torch::cat({features_reduced, features.index({torch::indexing::Slice(), indices_list[i] , torch::indexing::Slice(), torch::indexing::Slice()}).unsqueeze(1)}, 1);
  }
}