count_nonzero with requires grad (#50866)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/50792
fixes `count_nonzero` for tensors with requires_grad and also includes test
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50866
Reviewed By: ejguan
Differential Revision: D25996202
Pulled By: albanD
fbshipit-source-id: 61f2d7d62dd04e574a65ad03ef3a358b141fbae7