def test_gather_nd_int16_indices(self, compute_unit, backend, frontend, dynamic):
        """Test the indices access in torch model which gets lowered to gather_nd."""
        B, C, H, W, T = 1, 24, 64, 64, 32
        data = torch.rand(B, C, H, W)
        time = (torch.rand(1, T) * (C - 1)).to(torch.int)

        if frontend == TorchFrontend.TORCHSCRIPT:

            class DynamicModel(torch.nn.Module):
                def forward(self, data, time):
                    return data[torch.arange(B).unsqueeze(1), time, :, :]

            class StaticModel(torch.nn.Module):
                def forward(self, data):
                    return data[torch.arange(B).unsqueeze(1), time, :, :]

            torch_model = DynamicModel() if dynamic else StaticModel()
        else:

            class DynamicModel(torch.nn.Module):
                def __init__(self, B):
                    super().__init__()
                    self.slice0 = torch.arange(B).unsqueeze(1)

                def forward(self, data, time):
                    return data[self.slice0, time, :, :]

            class StaticModel(torch.nn.Module):
                def __init__(self, B, time):
                    super().__init__()
                    self.slice0 = torch.arange(B).unsqueeze(1)
                    self.time = time

                def forward(self, data):
                    return data[self.slice0, self.time, :, :]

            torch_model = DynamicModel(B) if dynamic else StaticModel(B, time)

        input_data = (data, time) if dynamic else data
        converter_input_type = [ct.TensorType(shape=data.shape)]
        if dynamic:
            converter_input_type.append(ct.TensorType(shape=time.shape, dtype=np.int32))

        mlmodel = self.run_compare_torch(
            input_data,
            torch_model,
            input_as_shape=False,
            converter_input_type=converter_input_type,
            frontend=frontend,
            backend=backend,
            compute_unit=compute_unit,
            minimum_deployment_target=ct.target.iOS17,
        )[1]
        gather_op = mlmodel._mil_program.find_ops(op_type="gather_nd")[0]
        assert gather_op.indices.dtype == types.int16 if dynamic else types.uint16
Test the indices access in torch model which gets lowered to gather_nd.
coremltools/converters/mil/frontend/torch/test/test_torch_ops.py