fix: support removing hook in the hook (#61250)
Summary:
Fixes: https://github.com/pytorch/pytorch/issues/58354
Problem:
Once a hook is called
https://github.com/pytorch/pytorch/blob/05c1e5b65571d2fa14b0f06c37ba970fd7cc43d0/torch/csrc/autograd/python_hook.cpp#L51-L54
If the hook has `handle.remove()` while executing and if there are no references to the hook function object then `python` is free to garbage collect.
At the subsequent call to
https://github.com/pytorch/pytorch/blob/05c1e5b65571d2fa14b0f06c37ba970fd7cc43d0/torch/csrc/autograd/python_hook.cpp#L54
we have `hook` pointing to invalid memory
Thus when we try to fetch the name for `hook` from `check_single_result` with
https://github.com/pytorch/pytorch/blob/05c1e5b65571d2fa14b0f06c37ba970fd7cc43d0/torch/csrc/autograd/python_hook.cpp#L175-L177
we get segfault.
Solution:
Temporarily increase the life-time of hook with `Py_INCREF` till we have verified the result.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61250
Reviewed By: iramazanli
Differential Revision: D29623826
Pulled By: soulitzer
fbshipit-source-id: c71322311f19066cafb7203980668868c59d4e5e
