How to select specific vector in 3D tensor beautifully?

There is a 3D tensor(batch_size * seq_len * hidden_dim), and I wanna get the different hidden vector for each sequence, do you have a beautiful method to do it except use for loop?

Tensor:
[[[1, 2, 3],
[4, 5, 6],
[7, 8, 9]],

[[11, 12, 13],
[14, 15, 16],
[17, 18, 19]]]

Target matrix:
[[4, 5, 6],
[17, 18, 19]]

and I have the index[1, 2] for select the specific vector, now I’m using for loop to get every vector for each sequence. And I found torch.gather and torch.index_select may be not suitable for my question.

Thank you for help!

1 Like

This should work:

x = torch.tensor([[[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]],
                  [[11, 12, 13],
                   [14, 15, 16],
                   [17, 18, 19]]])

idx = torch.tensor([1, 2])
x[torch.arange(x.size(0)), idx]
2 Likes

Thanks, I’ve tried slicing: x[:, idx], but it’s not correct. So using indexing is correct, thank you again.

thanks from 2022 XD, it works pretty well !