refactor autograd_hook (#54981)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54981
put part of codes in autograd_hook into functions, so that they can be used in the static graph training later on.
ghstack-source-id: 127755405
Test Plan: unit tests
Reviewed By: SciPioneer
Differential Revision: D27439508
fbshipit-source-id: a02a4b029841f5e7f11cfc5496bb7972ef53d878
