def test_fsdp_auto_wrap_policy_compute_sensitivity_call(
    mocker, num_sensitivity_workers, auto_wrap_policy
):
    """
    Test compute_sensitivity passes fsdp_auto_wrap_policy argument correctly to
    impl methods
    """
    model = torch.nn.Linear(5, 10)

    mocker.patch("coremltools.optimize.torch.palettization.sensitive_k_means._torch.save")
    mocker.patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._torch.load",
        Mock(return_value=model.state_dict()),
    )
    mocker.patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._torch.cuda.is_available",
        Mock(return_value=True),
    )
    mock_ctx = Mock()
    mocker.patch(
        "coremltools.optimize.torch.palettization.sensitive_k_means._mp.get_context",
        Mock(return_value=mock_ctx),
    )
    mock_compute_sen_single_worker = Mock()
    mocker.patch.object(
        SKMPalettizer,
        "_compute_sensitivity_impl_single_worker",
        mock_compute_sen_single_worker,
    )
    mock_dataset = Mock()
    mocker.patch.object(SKMPalettizer, "_get_dataset", Mock(return_value=mock_dataset))
    mocker.patch.object(SKMPalettizer, "_process_sensitivity")

    palettizer = SKMPalettizer(model)

    dataloader = []
    loss_fn = lambda mod, dat: mod(dat)

    palettizer.compute_sensitivity(
        dataloader=dataloader,
        loss_fn=loss_fn,
        sensitivity_path=None,
        num_sensitivity_workers=num_sensitivity_workers,
        fsdp_auto_wrap_policy=auto_wrap_policy,
    )

    if num_sensitivity_workers > 1:
        for rank in range(num_sensitivity_workers):
            mock_ctx.Process.assert_any_call(
                target=palettizer._compute_sensitivity_impl_multiple_workers,
                args=(
                    rank,
                    num_sensitivity_workers,
                    mock_dataset,
                    loss_fn,
                    None,
                    auto_wrap_policy,
                ),
                name=f"Process-{rank}",
            )
    else:
        mock_compute_sen_single_worker.assert_called_once_with(mock_dataset, loss_fn, None)

    Test compute_sensitivity passes fsdp_auto_wrap_policy argument correctly to
    impl methods
    
coremltools/test/optimize/torch/palettization/test_sensitive_k_means.py