fix: pass train_bn to freeze in BackboneFinetuning.freeze_before_training

freeze_before_training() called self.freeze(pl_module.backbone)
without forwarding self.train_bn, so BatchNorm layers stayed
trainable during the initial frozen phase regardless of the
train_bn setting. The unfreezing path already passed the flag
correctly.

Author: javierdejesusda

src/lightning/pytorch/callbacks/finetuning.py

@@ -455,7 +455,7 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
 
     @override
     def freeze_before_training(self, pl_module: "pl.LightningModule") -> None:
-        self.freeze(pl_module.backbone)
+        self.freeze(pl_module.backbone, train_bn=self.train_bn)
 
     @override
     def finetune_function(self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer) -> None:

tests/tests_pytorch/callbacks/test_finetuning_callback.py

@@ -73,6 +73,37 @@ def train_dataloader(self):
     assert model.backbone.has_been_used
 
 
+def test_finetuning_callback_train_bn_false(tmp_path):
+    """Test that BackboneFinetuning respects train_bn=False during the initial freeze phase."""
+    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 = nn.Linear(32, 2)
+
+        def forward(self, x):
+            return self.layer(self.backbone(x))
+
+        def configure_optimizers(self):
+            return torch.optim.SGD(self.layer.parameters(), lr=0.1)
+
+        def train_dataloader(self):
+            return DataLoader(RandomDataset(32, 64), batch_size=2)
+
+    model = FinetuningBoringModel()
+    callback = BackboneFinetuning(unfreeze_backbone_at_epoch=3, train_bn=False, verbose=False)
+
+    trainer = Trainer(limit_train_batches=4, default_root_dir=tmp_path, callbacks=[callback], max_epochs=1)
+    trainer.fit(model)
+
+    # With train_bn=False, BatchNorm should be fully frozen (not trainable, no running stats)
+    assert not model.backbone[1].weight.requires_grad
+    assert not model.backbone[1].bias.requires_grad
+    assert not model.backbone[1].track_running_stats
+
+
 class TestBackboneFinetuningWarningCallback(BackboneFinetuning):
     def finetune_function(self, pl_module, epoch: int, optimizer):
         """Called when the epoch begins."""