Thanks !
My idea was this (not that scatter and scatter_add only exist as inplace ops but you can clone vocab_dist before the scatter_ if you don’t want to change it inplace):
vocab_dist  = output_layer(hiddens)  # output layer is linear layer and output size: [batch_size, 50000]
scores = torch.randn(16, 400)   # [batch size, number of time steps]
indices = torch.randint(0, 50000, (16, 400))   # [batch_size, number of steps]
current_values = vocab_dist.gather(1, indices)
new_val = torch.max(current_values, scores)
vocab_dist.scatter_(1, indices, new_val)