Loosen output shape restrictions on GPT-style models (#25188)
* Loosen output shape restrictions on GPT-style models
* Use more self-explanatory variables
* Revert "Use more self-explanatory variables"
This reverts commit 5fd9ab39119558b7e750f61aa4a19014dccc5ed5.