Suppose I have a list of indices: ind = torch.tensor([0, 2, 1, 3]), and a 4x4 tensor A.
I need to assign some values to the items A[i, ind[i], for i=1,4. For example:
for i in range(4):
    A[i, ind[i]] = r[i]

where r contains the target values in this example.
Given A, ind, r, Is there a way to vectorize this operation?