def test_index_cast_to_int(self, compute_unit, backend):
        """Test the index are cast into the correct dtype."""

        class IndexModel(torch.nn.Module):
            def forward(self, x, y):
                mask = y > 2
                index = mask.sum(0).unsqueeze(0)
                return x[index, :]

        x = torch.rand((2, 3))
        y = torch.Tensor([1.0, 2.0, 3.0])

        self.run_compare_torch(
            [x, y],
            IndexModel(),
            backend=backend,
            compute_unit=compute_unit,
            input_as_shape=False,
            minimum_deployment_target=ct.target.iOS17,
        )
Test the index are cast into the correct dtype.
coremltools/converters/mil/frontend/torch/test/test_torch_ops.py