Refactor fast gradcheck (#55871)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55871
Test Plan: Imported from OSS
Reviewed By: albanD
Differential Revision: D28096549
Pulled By: soulitzer
fbshipit-source-id: ee8b71fbd03ee581e71cdfcfd5e2258adefe15a6
