Hello, say I have a tensor
c of 3D coordinates with shape (B x P x 3), where B is the batch size, and another one with indexes
idx and shape (B x I x 2):
c = torch.tensor([[[1,1,2],[5,2,3],[6,1,8]]]) idx = torch.tensor([[[0,1],[1,2]]])
(Here the batch size is 1, just for simplicity).
I want to slice the coordinates according to the idx tensor, so the result should be a tensor with shape (B x I x P x 3) and values
result[:,0] = torch.tensor([[1,1,2],[5,2,3]]) result[:,1] = torch.tensor([[5,2,3],[6,1,8]])
I have tried using slicing
c[:,idx], but that only works apparently when B=1, although it adds dimensions I don’t want. I also tried with
torch.gather(c, dim=1, index=idx). What is the way to do this?