Tensor indexing from image feature batch using tensor of indices

vision
1

Hi everyone. I have an image feature tensor of shape (16, 500, 500, 64) and a tensor of indices of shape (16, 6890, 2) which are essentially pixel locations I am meant to slice. How do I do this in PyTorch without using for loops (since I am doing this inside the model of my network and loops being slow doesn’t seem like a great idea). Any help is appreciated. Thanks in advance!

Just to make things clear, the slicing is supposed to give an output of size (16, 6890, 64), i.e. feature slices corresponding to pixel locations.

1 Like
3

Found a possible solution here : How to do the tf.gather_nd in pytorch? - #18 by Cogito2012. On applying this, I do get an output of desired shape. However, this error keeps triggered for the following operations ( RuntimeError: CUDA error: device-side assert triggered)

Edit : This works! Turns out there was something wrong with my indices.