https://github.com/lucidrains/PaLM-rlhf-pytorch/blob/6b02ee329106baff78e293afa7d1d2e6dd4e5ca2/palm_rlhf_pytorch/utils.py#L60 Using the sorted indices to index the sorted indices does not make sense. I think it may be `return logits.scatter(1, sorted_indices, sorted_logits)`
PaLM-rlhf-pytorch/palm_rlhf_pytorch/utils.py
Line 60 in 6b02ee3
Using the sorted indices to index the sorted indices does not make sense.
I think it may be
return logits.scatter(1, sorted_indices, sorted_logits)