def test_finetuning_callback_warning(tmp_path):
    """Test finetuning callbacks works as expected."""
    seed_everything(42)

    class FinetuningBoringModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.backbone = nn.Linear(32, 2, bias=False)
            self.layer = None
            self.backbone.has_been_used = False

        def forward(self, x):
            self.backbone.has_been_used = True
            return self.backbone(x)

        def train_dataloader(self):
            return DataLoader(RandomDataset(32, 64), batch_size=2)

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

    chk = ModelCheckpoint(dirpath=tmp_path, save_last=True)

    model = FinetuningBoringModel()
    model.validation_step = None
    callback = TestBackboneFinetuningWarningCallback(unfreeze_backbone_at_epoch=3, verbose=False)

    with pytest.warns(UserWarning, match="Did you init your optimizer in"):
        trainer = Trainer(limit_train_batches=1, default_root_dir=tmp_path, callbacks=[callback, chk], max_epochs=2)
        trainer.fit(model)

    assert model.backbone.has_been_used
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=3)
    trainer.fit(model, ckpt_path=chk.last_model_path)
Test finetuning callbacks works as expected.
tests/tests_pytorch/callbacks/test_finetuning_callback.py