MPS on mac Book pro: Converting Bool tensor to float

I am trying to train NLP model on 16" macbook pro 2022,
OS 12.5 with M1 Pro Chip
pytorch version: 1.13.0
platform.platform() → ‘macOS-12.5-arm64-arm-64bit’
Running following code

import torch
lengths=torch.randint(8, 19, (10,)).to('mps')
max_len=lengths.data.max()
batch_size=lengths.size(0)
seq_range = torch.arange(0, max_len).long()
seq_range_exp = seq_range.unsqueeze(0).expand(batch_size, max_len)
seq_range_exp=seq_range_exp.to(lengths)
seq_length_exp = lengths.unsqueeze(1).expand_as(seq_range_exp)
length_mask=seq_range_exp < seq_length_exp
print(length_mask.float())
print(length_mask[:, 1].float())

Output:

tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 0.]],
device=‘mps:0’)
tensor([1.0000e+00, 2.3694e-38, 2.3694e-38, 3.6013e-43, 2.3694e-38, 2.3694e-38,
3.6013e-43, 1.4013e-45, 2.3694e-38, 2.3694e-38], device=‘mps:0’)

/Users/ngupta/opt/anaconda3/envs/ml/lib/python3.9/site-packages/torch/_tensor_str.py:115: UserWarning: The operator ‘aten::nonzero’ is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
nonzero_finite_vals = torch.masked_select(

print(length_mask[:, 1].float()) should print column 1 of the tensor, which is all 1.0
obviously
“tensor([1.0000e+00, 2.3694e-38, 2.3694e-38, 3.6013e-43, 2.3694e-38, 2.3694e-38,
3.6013e-43, 1.4013e-45, 2.3694e-38, 2.3694e-38], device=‘mps:0’)”
is not correct. I will appreciate an explanation and a remedy to correct this.