Select elements of Tensor based on index tensor along same dimensions
I have the following two tensors
input shape: 16 32 32 3
index shape: 16 32 32 2
output shape: 16 32 32 3

The formula for the output would be:
output[b, h, w] = input[b, index[b, h, w, 0], index[b, h, w, 1]]

I tried to use torch.gather but I was not able to formulate the previous assignment.
Does anyone know how to do this in an efficient manner? Thanks!
For context: input contains a batch of 16 elemens where each one is a tensor of 32x32 that containts 3D points. index is a mapping from position to 3D point.

You can achieve this by unraveling the indices from index (on dimensions 1 and 2) in order to index input on a single dimension using torch.gather.
This requires to expand the shape of the indexer to fit the shape of input:
Here is an example with some dummy data:
&gt;&gt;&gt; x = torch.rand(16, 32, 32, 3)
&gt;&gt;&gt; index = torch.randint(0, 10, (16,32,32,2))

Some manipulation on index is required to unravel the values:
&gt;&gt;&gt; unraveled = x.size(1)*index[..., 0] + index[..., 1]
&gt;&gt;&gt; u = unraveled.flatten(1).unsqueeze(-1).expand(-1, -1, x.size(-1))

Now u, reshaped expanded version of index has a shape of (16, 1024, 3).
The indexed tensor also needs to be flattened:
&gt;&gt;&gt; x.flatten(1, 2)
torch.Size([16, 1024, 3])

Finally, you can gather on dim=1 (keep in mind the result needs to be reshaped to the desired shape i.e. the input's shape):
&gt;&gt;&gt; out = input.flatten(1,2).gather(1, u).reshape_as(x)

