How does one vectorize partial index selection in pytorch?

1

I want to change the values of entries of my MxN matrix but I have which column I want for each row on a vector. I tried:

scores[:column_indices] = 0

but that sets everything to zero instead of for all row change column indicated by column_indices to zero. I guess in python loops it would be:

for k,i in enumerate(column_indices):
    scores[k,i] = 0

does someone know how to vectorize that code? I feel it should be easy but I’m having issues doing it.

2

Could you print the column_indices?

3

Yes of course! :slight_smile:

tensor([ 6,  7,  9,  ...,  8,  5,  6], device='cuda:0')
4

Haha, sorry, didn’t see the enumerate and thought column_indices is two dimensional. :wink:

However, I didn’t come up with a faster solution. :frowning:

1 Like
5

haha, no worries :slight_smile:

perhaps this is not how numpy or matlab indices work and expected something out of the ordinary in pytorch…?

6

Well, there are some options like using narrow or perhaps scatter will also work.
However, I don’t know how to create the indices without a for loop, which makes my approaches slower than your loop.

1 Like