def test_base_finetuning_internal_optimizer_metadata(tmp_path):
    """Test the param_groups updates are properly saved within the internal state of the BaseFinetuning Callbacks."""

    seed_everything(42)

    class FreezeModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.layer = nn.Sequential(
                nn.Linear(32, 32, bias=False),
                nn.Linear(32, 32, bias=True),
                nn.Linear(32, 32, bias=False),
                nn.Linear(32, 32, bias=True),
                nn.Linear(32, 32, bias=False),
                nn.Linear(32, 2, bias=True),
            )

        def forward(self, x):
            return self.layer(x)

        def configure_optimizers(self):
            return torch.optim.SGD(self.layer[0].parameters(), lr=0.1)

    cb = OnEpochLayerFinetuning()
    chk = ModelCheckpoint(dirpath=tmp_path, save_last=True)
    model = FreezeModel()
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=5, limit_train_batches=1, callbacks=[cb, chk])
    trainer.fit(model)
    assert len(cb._internal_optimizer_metadata[0]) == 6
    assert cb._internal_optimizer_metadata[0][0]["params"] == ["layer.0.weight"]
    assert cb._internal_optimizer_metadata[0][1]["params"] == ["layer.1.weight", "layer.1.bias"]
    assert cb._internal_optimizer_metadata[0][2]["params"] == ["layer.2.weight"]
    assert cb._internal_optimizer_metadata[0][3]["params"] == ["layer.3.weight", "layer.3.bias"]
    assert cb._internal_optimizer_metadata[0][4]["params"] == ["layer.4.weight"]
    assert cb._internal_optimizer_metadata[0][5]["params"] == ["layer.5.weight", "layer.5.bias"]

    model = FreezeModel()
    cb = OnEpochLayerFinetuning()
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=10, callbacks=[cb])
    with pytest.raises(IndexError, match="index 6 is out of range"):
        trainer.fit(model, ckpt_path=chk.last_model_path)
Test the param_groups updates are properly saved within the internal state of the BaseFinetuning Callbacks.
tests/tests_pytorch/callbacks/test_finetuning_callback.py