Improve support for backwards hooks (#119525)
Summary:
X-link: https://github.com/pytorch/pytorch/pull/119525
Approved by: https://github.com/yanboliang, https://github.com/anijain2305
Reviewed By: huydhn
Differential Revision: D53634801
Pulled By: jansel
fbshipit-source-id: d18819e331cd408f099ef5101cdbe9bc3d1ae172
