Tensor assigment when doubly indexed

I’m having trouble understanding why this piece of code does not work as expected:

a = torch.rand(1000)
idx1 = torch.arange(500, 600)
idx2 = torch.rand(100) < 0.5

a[idx1][idx2] = 1.0

I expect the result to be all ones, but the tensor is unchanged. If instead, I index only once like so:

a[idx1] = 1.0

Then, as expected I get all ones. Why does the double index not work and how can I achieve this behavior?

What you are loking for is:

a[idx1[idx2]] = 1.0

(Mask idx1 indexes tensor using boolean idx2 tensor before slicing tensor a)

Maybe my example wasn’t broad enough. Your solution doesn’t work if the idx shapes don’t match. Say idx1 is a boolean array of shape 1000, and idx2 is of shape idx1.sum(). What would you do then?

Something like torch.where(idx1)[0][idx2] could work, but that seems very inefficient.