def test_permanent_when_model_is_saved_multiple_times(
    tmp_path, caplog, prune_on_train_epoch_end, save_on_train_epoch_end
):
    """When a model is saved multiple times and make_permanent=True, we need to make sure a copy is pruned and not the
    trained model if we want to continue with the same pruning buffers."""
    if prune_on_train_epoch_end and save_on_train_epoch_end:
        pytest.xfail(
            "Pruning sets the `grad_fn` of the parameters so we can't save"
            " right after as pruning has not been made permanent"
        )

    class TestPruning(ModelPruning):
        def on_save_checkpoint(self, trainer, pl_module, checkpoint):
            had_buffers = hasattr(pl_module.layer.mlp_3, "weight_orig")
            super().on_save_checkpoint(trainer, pl_module, checkpoint)
            assert "layer.mlp_3.weight_orig" not in checkpoint["state_dict"]
            if had_buffers:
                assert hasattr(pl_module.layer.mlp_3, "weight_orig")

    model = TestModel()
    pruning_callback = TestPruning(
        "random_unstructured",
        parameters_to_prune=[(model.layer.mlp_3, "weight")],
        verbose=1,
        make_pruning_permanent=True,
        prune_on_train_epoch_end=prune_on_train_epoch_end,
    )
    ckpt_callback = ModelCheckpoint(
        monitor="test", save_top_k=2, save_last=True, save_on_train_epoch_end=save_on_train_epoch_end
    )
    trainer = Trainer(
        default_root_dir=tmp_path,
        logger=False,
        callbacks=[pruning_callback, ckpt_callback],
        max_epochs=3,
        enable_progress_bar=False,
    )
    with caplog.at_level(INFO):
        trainer.fit(model)

    actual = [m.strip() for m in caplog.messages]
    actual = [m for m in actual if m.startswith("Applied")]
    percentage = r"\(\d+(?:\.\d+)?%\)"
    expected = [
        rf"Applied `RandomUnstructured`. Pruned: \d+\/66 {percentage} -> \d+\/66 {percentage}",
        rf"Applied `RandomUnstructured`. Pruned: \d+\/66 {percentage} -> \d+\/66 {percentage}",
        rf"Applied `RandomUnstructured`. Pruned: \d+\/66 {percentage} -> \d+\/66 {percentage}",
    ]
    expected = [re.compile(s) for s in expected]
    assert all(regex.match(s) for s, regex in zip(actual, expected))

    # removed on_train_end
    assert not hasattr(model.layer.mlp_3, "weight_orig")

    model = TestModel.load_from_checkpoint(trainer.checkpoint_callback.kth_best_model_path)
    assert not hasattr(model.layer.mlp_3, "weight_orig")
    model = TestModel.load_from_checkpoint(trainer.checkpoint_callback.last_model_path)
    assert not hasattr(model.layer.mlp_3, "weight_orig")
When a model is saved multiple times and make_permanent=True, we need to make sure a copy is pruned and not the
    trained model if we want to continue with the same pruning buffers.
tests/tests_pytorch/callbacks/test_pruning.py