How to auto-regression generate？

[yileld]

In eval_needle.py, it gather( ) and undo_extract_local( ) the preds to get the whole preds, then get pred token by prompt_length. In auto-regression mode, I just need next token, can I just get the pred token without gather( ) and undo_extract_local( )?