def test_multinomial_probs_instead_of_logits(self, compute_unit, backend):
        """
        Verify the input to multinomial is probs instead of logits.

        When the number of drawing is large, the drawing results could tell us if the input is probs
        or logits. In this test we use only 2 classes, so we can compare the number of `1` in results
        to verify if the input is taken a logarithm or not.
        """

        class TestModel(nn.Module):
            def forward(self, x):
                return torch.multinomial(x, 1000, replacement=True)

        input_data = torch.tensor([0.01, 0.1], dtype=torch.float)
        torch_model = TestModel()
        torch_model.eval()
        traced_model = torch.jit.trace(torch_model, input_data)
        mlmodel = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=input_data.shape, dtype=np.float16)],
            outputs=[ct.TensorType(name="output", dtype=np.float16)],
            convert_to="mlprogram",
            compute_units=ct.ComputeUnit.CPU_ONLY,
            minimum_deployment_target=ct.target.iOS16,
        )

        if ct.utils._is_macos():
            mlmodel_out = mlmodel.predict({"input": input_data.numpy()})["output"]
            torch_out = torch_model(input_data).numpy()
            # The counting of 1 in PyTorch and CoreML output should be similar.
            assert np.abs(np.sum(mlmodel_out) - np.sum(torch_out)) / mlmodel_out.size < 0.05

        Verify the input to multinomial is probs instead of logits.

        When the number of drawing is large, the drawing results could tell us if the input is probs
        or logits. In this test we use only 2 classes, so we can compare the number of `1` in results
        to verify if the input is taken a logarithm or not.
        
coremltools/converters/mil/frontend/torch/test/test_torch_ops.py