How to do the tf.gather_nd in pytorch?

3

Yes, there are equivalent operations in pytorch. Try something like the following:

Simple indexing into matrix:

x = torch.randn(2, 2)
indices = torch.ByteTensor([[0, 0],[1,1]])
x.masked_select(indices)

Slice indexing into matrix:

x = torch.randn(2, 2)
indices = torch.LongTensor([1, 0])
x.index_select(0, indices)

Indexing into a 3-tensor:

x = torch.randn(1, 4, 2)
x[:,:,1]

Batched indexing into a matrix:

x = torch.randn(2, 2)
indices = torch.LongTensor([[0, 0],[1,1]])
[x[i] for i in indices]
2 Likes