[MPS] Calculate nonzero count first before running nonzero op (#102052)
Summary of changes:
- Calculate nonzero count first before running nonzero op
- allocate only 1 element when calling .item(), and blit only the size of destination
Pull Request resolved: https://github.com/pytorch/pytorch/pull/102052
Approved by: https://github.com/kulinseth
