Move more stuff to after_aot (#99557)
Not sure why this didn't work first time around. Second time's a charm.
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/99557
Approved by: https://github.com/anijain2305
