Use guard_manager consistently instead of check_fn (#138896)
Summary:
X-link: https://github.com/pytorch/pytorch/pull/138896
Approved by: https://github.com/williamwen42, https://github.com/jansel
ghstack dependencies: #138512
Reviewed By: wdvr
Differential Revision: D65030963
Pulled By: anijain2305
fbshipit-source-id: 7423473e4c3613aea42e13a64eae9c417c876964
