Hi All,

I have two 3D tensors X and Q of shape (5, 16, 128) on which I do cosine similarity on 2nd dim to get a (5, 16) cosine-similarity vector. I then sort this cosine-similarity vector, to get indices of most-to-least similar vectors in the original vector Q. In code

```
X = torch.ones((5, 16, 128))
Q = torch.randn(5, 16, 128) # dim=2 is features
neg_samples = 11
pos_samples = 5
cos_sim = torch.nn.CosineSimilarity(dim=2, eps=1e-6)
sim_scores = cos_sim(X, Q) # sim_scores output shape: [5,16]
_, sorted_idxs = torch.topk(sim_scores, k=sim_scores.shape[1], sorted=True, dim=1) #sort sim_scores to get sorted indexes same shape [5, 16]
#since sorted_idxs is sorted descending on similarity, we can slice to divide pos/neg samples.
pos_idxs = sorted_idxs[:, :pos_samples] #gives [5, 5] most positive idxs from sorted_idxs
neg_idxs = sorted_idxs[:, pos_samples:(pos_samples+neg_samples)] #gives [5, 11] most negative idxs from sorted_idxs
```

How do i retrieve the corresponding tensors from `Q`

which is of shape `[5, 16, 128]`

using these `pos_idxs`

shape `[5,5]`

and `neg_idxs`

shape `[5,11]`

?

I tried `Q[ pos_idxs, :]`

but i get output tensor of shape `[5, 5, 16, 128]`

. I want to slice `Q`

into `[5, 5, 128]`

and `[5, 11, 128]`

two separate tensors based on 2D `pos_idxs`

and `neg_idxs`

index tensors.

How do i do this in pytorch using tensor operations?

thanks