Fix test_hook_* issues (#40135)
Summary:
Follows https://github.com/pytorch/pytorch/issues/38972
Some of the changes asked by albanD in the above review are appliable to the regular hooks tests.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40135
Differential Revision: D22091389
Pulled By: albanD
fbshipit-source-id: e1004213276bfb189167b9870e1a88b3d23b458c
