Fix MPS test sanity (#84889)
Follow up after https://github.com/pytorch/pytorch/pull/84834
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84889
Approved by: https://github.com/tugsbayasgalan, https://github.com/janeyx99, https://github.com/ZainRizvi
