How does one vectorize partial index selection in pytorch?

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.

Could you print the column_indices?

Yes of course! :slight_smile:

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

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

haha, no worries :slight_smile:

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

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