def test_index_int_index_case_6(
        self, compute_unit, backend, frontend, input_dtype, shape, minimum_deployment_target
    ):
        """Only one axis is sliced + nd mode."""
        class IndexModel(torch.nn.Module):
            def forward(self, x):
                if len(shape) == 2:
                    index = torch.tensor([0, 0, 0, 0, 0, 0])
                    index = index.view(2, 3)
                    return x[index, :]
                elif len(shape) == 4:
                    index = torch.tensor([0, 1, 2, 3, 0, 1])
                    index = index.view(3, 2)
                    return x[:, index]

        model = IndexModel()
        self.run_compare_torch(
            shape,
            model,
            frontend=frontend,
            backend=backend,
            compute_unit=compute_unit,
            rand_range=(0, 2),
            input_dtype=input_dtype,
            minimum_deployment_target=minimum_deployment_target,
        )
Only one axis is sliced + nd mode.
coremltools/converters/mil/frontend/torch/test/test_torch_ops.py