C++ API tensor indexing

I’m trying to replicate the following Python code in the PyTorch C++ frontend:

value = torch::rand({2, 2, 4});
value = value[..., :1]
# value.shape == torch.Size([2, 2, 1])

But I can’t for the life of me get indexing to work, and I can’t seem to find any documentation on indexing in C++.

Here’s what I think should work:

long two_array[2] = {0, 1};
auto two = torch::from_blob(two_array, {2}, torch::kLong);
long one_array[1] = {0};
auto one = torch::from_blob(one_array, {2}, torch::kLong);
value = value.index({two, two, one});

However that gives the following error

invalid argument 4: Index tensor must have same dimensions as input tensor at /home/px046/dev/pytorch/aten/src/TH/generic/THTensorEvenMoreMath.cpp:438

How do you do indexing like this in the C++ frontend?

1 Like

narrow is a good fit for getting a slice.

Best regards

Thomas

1 Like

value = value.narrow(-1, 0, 1) worked great, thanks!