index_put failing when indices are bool type

## 🐞Describing the bug
index_put seems to fail when I have boolean indices . Please see following steps to reproduce

## To Reproduce
- Please add a minimal code example that can reproduce the error when running it.
```python
import torch
from torch import nn
import coremltools as ct

class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.not_a_point_embed = torch.ones(1, 2)
self.point_embed = torch.ones(1, 2)*2

@torch.no_grad()
def forward(self, point_embedding, labels):
point_embedding[labels == 0] += self.not_a_point_embed
point_embedding[labels == 1] += self.point_embed
return point_embedding

model = SimpleNet().eval()
point_embedding = torch.zeros(1, 3, 2)
labels = torch.tensor([[1, 0, 1]])

traced_model = torch.jit.trace(model, (point_embedding, labels))

mlmodel=ct.convert(traced_model,
convert_to="mlprogram",
inputs=[
ct.TensorType(name="point_embedding", shape=point_embedding.shape),
ct.TensorType(name="labels", shape=labels.shape)])
```

## System environment (please complete the following information):
- coremltools version: coremltools installed from `main` branch of the repo
- OS (e.g. MacOS version or Linux type): MacOS
- Any other relevant version information (e.g. PyTorch or TensorFlow version): Pytorch version `Torch version 2.2.0`

## Additional context
Its failing with assert in index_put torch_op:

```
@register_torch_op
def index_put(context, node):
inputs = _get_inputs(context, node, expected=4)
x = inputs[0]
indices = inputs[1]
values = inputs[2]
accumulate = inputs[3].val
rank = x.rank
mode = "add" if accumulate else "update"

indices_type = indices[0].sym_type.get_primitive()

if types.is_bool(indices_type):
assert len(indices) == 1, "Unsupported index_put_ usage."
indices = indices[0]
assert (
indices.shape == x.shape
), "indices shape must equal to input shape for index put operation."
indices = mb.cast(x=indices, dtype="int32")
indices = mb.non_zero(x=indices)
```

It looks like assert is failing in the shape-check here:

```
assert (
indices.shape == x.shape
), "indices shape must equal to input shape for index put operation."
```

Is this check correct?
I tried changing it to `indices.shape == x.shape[-1]`, but `scatter_nd` fails subsequently in further checks.

Any ideas on what's going on here?
Appreciate any help in further troubleshooting this.

Start by running the provided PyTorch and coremltools conversion example, then inspect the index_put torch operation implementation and its boolean-index shape handling. Done means the model converts successfully with boolean indices and the failing shape assertion is resolved without causing the later scatter_nd checks to fail.

Written by the indexing model from the issue text.