Hi,
Just wanted to link this Big generate() refactor - 
Transformers - Hugging Face Forums for people who are looking into this in the future! I recently required gradients computed with respect to the logits but was unable to do so until I found the above link.
This discussion: Question about greedy_search - Transformers - Hugging Face Forums was also useful and provided a more concrete example to the above.
Thank you.