Revert "Remove extra graph breaks (#118987)"
This reverts commit 9a8e3b07d75e3e9bb902f81b4b6e1042bbe06b58.
Reverted https://github.com/pytorch/pytorch/pull/118987 on behalf of https://github.com/eellison due to reverting because it causes regression ([comment](https://github.com/pytorch/pytorch/pull/118987#issuecomment-1928224447))
