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)