def logprobs_from_logits(logits, labels):
    """
    See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591
    """
    logp = F.log_softmax(logits, dim=2)
    logpy = torch.gather(logp, 2, labels.unsqueeze(2)).squeeze(-1)
    return logpy

    See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591
    
promptist/trlx/trlx/utils/modeling.py