def test_callbacks_restore(tmp_path):
    """Test callbacks restore is called after optimizers have been re-created but before optimizer states reload."""
    chk = ModelCheckpoint(dirpath=tmp_path, save_last=True)

    model = FinetuningBoringModel()
    callback = TestCallbacksRestoreCallback()

    trainer_kwargs = {
        "default_root_dir": tmp_path,
        "limit_train_batches": 1,
        "limit_val_batches": 1,
        "callbacks": [callback, chk],
        "max_epochs": 2,
    }

    trainer = Trainer(**trainer_kwargs)
    trainer.fit(model)

    # only 1 optimizer
    assert len(callback._internal_optimizer_metadata) == 1

    # only 2 param groups
    assert len(callback._internal_optimizer_metadata[0]) == 2

    # original parameters
    expected = {
        "lr": 0.1,
        "momentum": 0,
        "dampening": 0,
        "weight_decay": 0,
        "nesterov": False,
        "params": ["layer.3.weight", "layer.3.bias"],
        "maximize": False,
        "foreach": None,
        "differentiable": False,
    }
    if _TORCH_GREATER_EQUAL_2_3:
        expected["fused"] = None

    assert callback._internal_optimizer_metadata[0][0] == expected

    # new param group
    expected = {
        "lr": 0.01,
        "momentum": 0,
        "dampening": 0,
        "weight_decay": 0,
        "nesterov": False,
        "params": ["layer.0.weight", "layer.0.bias"],
        "maximize": False,
        "foreach": None,
        "differentiable": False,
    }
    if _TORCH_GREATER_EQUAL_2_3:
        expected["fused"] = None

    assert callback._internal_optimizer_metadata[0][1] == expected

    trainer_kwargs["max_epochs"] = 3

    trainer = Trainer(**trainer_kwargs)
    trainer.fit(model, ckpt_path=chk.last_model_path)
Test callbacks restore is called after optimizers have been re-created but before optimizer states reload.
tests/tests_pytorch/callbacks/test_finetuning_callback.py