🐛 Describe the bug

The output of masked.log_softmax on CPU with inputs of cpu_args: [tensor(-8.4844, dtype=torch.float16, requires_grad=True), 0] and cpu_kwargs: {'mask': tensor(False)}, i.e., the masked element, is nan. While the result on MPS is 0. https://github.com/pytorch/pytorch/actions/runs/6542102776/job/17764726988.
masked.softmax and masked.softmin should have the same issue.

Versions

PyTorch: main branch.

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev @ezyang @gchanan @zou3519 @kadeng