Hi Umbriel!



 umbriel:

Result should be: [[X, 2, 1], [2, 4, Y], [6, Z, 2]]


You could use indexing with assignment:
>>> import torch
>>> torch.__version__
'1.10.2'
>>> A = torch.tensor ([[3, 2, 1], [2, 4, 5], [6, 3, 2]])
>>> B = torch.tensor ([100, 200, 300])
>>> Index = torch.tensor ([0, 2, 1])
>>> A[torch.arange (3), Index] = B
>>> A
tensor([[100,   2,   1],
        [  2,   4, 200],
        [  6, 300,   2]])

Best.
K. Frank