Pytorch изменить массив со списком индексов

Question

Предположим, у меня есть список индексов и я хочу изменить существующий массив с этим списком. В настоящее время единственный способ сделать это - использовать цикл for следующим образом. Просто интересно, есть ли более быстрый / эффективный способ.

torch.manual_seed(0)
a = torch.randn(5,3)
idx = torch.Tensor([[1,2], [3,2]], dtype=torch.long)
for i,j in idx:
    a[i,j] = 1

Я изначально предполагал, что gather или же index_select мог бы как-то ответить на этот вопрос, но, глядя на документацию, это не похоже на ответ.

В моем конкретном случае a - это 5-мерный вектор, а idx - это вектор Nx5. Таким образом, выход (после подписки с чем-то вроде a[idx]) Я ожидаю, что это (N,) вектор

Ответ

Спасибо @shai ниже, ответ, который я искал, был:a[idx.t().chunk(chunks=2,dim=0)], Взято из этого ТАКОГО ответа.

Answer 1

Это довольно просто

a[idx[:,0], idx[:,1]] = 1

Вы можете найти более общее решение в этой теме.