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?