Yeah, I’m not sure either if there is one clean solution using .gather.
I tried to use multiple calls to gather:
x.gather(1, idx_h).gather(2, idx_w)

, but obviously the second gather call will index x as x[i][j][idx_w[i][j][k]], which is not what we want.
There is one approach using strides, but it’s not really straightforward to understand, what is happening.
I’m really not a fan of this code and I think there might be a better approach.
However, for the sake of completeness:
idx0 = torch.arange(N).view(N, 1, 1).repeat(1, H, W).view(-1)
x.view(-1)[idx0 * x.stride()[0] + idx_h.view(-1) * x.stride(1) + idx_w.view(-1)].view(N, H, W)