Use safe_is_leaf to test leafness (#102706)
This fixes one of the problems in https://github.com/pytorch/pytorch/issues/101160#issuecomment-1570376548
but I don't have a test case because the full example is fairly
difficult to minify.
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/102706
Approved by: https://github.com/bdhirsh
