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

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

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

        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.7)
            return [optimizer], [lr_scheduler]

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

    model = FinetuningBoringModel()
    callback = TestBackboneFinetuningCallback(unfreeze_backbone_at_epoch=3, verbose=False)

    trainer = Trainer(limit_train_batches=4, default_root_dir=tmp_path, callbacks=[callback], max_epochs=8)
    trainer.fit(model)

    assert model.backbone.has_been_used
Test finetuning callbacks works as expected.
tests/tests_pytorch/callbacks/test_finetuning_callback.py