Adding values to a tensor at specified indices in a dimension

Not sure if this has been asked before.

I have a tensor, a of size bsz * 10000
a tensor, idx of size bsz * 30 which contains long values - index values lying between 0 and 10000
a tensor, b of size bsz * 30 float values

how can I add values in b to tensor a along the dimension 2 (10000) at indices specified by idx ?

I’m not sure, if you would like to index b also using idx, but if that’s the case, you could try this code:

batch_size = 5
nb_features = 10
a = torch.zeros(batch_size, nb_features)
idx = torch.empty(batch_size, 3, dtype=torch.long).random_(nb_features)
b = torch.ones(batch_size, 3)
a.scatter_(1, idx, b)