Fix train_bn arg

pl_examples/domain_templates/computer_vision_fine_tuning.py

@@ -71,7 +71,7 @@ def __init__(self, milestones: tuple = (5, 10), train_bn: bool = False):
         self.train_bn = train_bn
 
     def freeze_before_training(self, pl_module: pl.LightningModule):
-        self.freeze(modules=pl_module.feature_extractor, train_bn=False)
+        self.freeze(modules=pl_module.feature_extractor, train_bn=self.train_bn)
 
     def finetune_function(self, pl_module: pl.LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
         if epoch == self.milestones[0]:
@@ -164,21 +164,19 @@ class TransferLearningModel(pl.LightningModule):
     def __init__(
         self,
         backbone: str = "resnet50",
-        train_bn: bool = True,
         milestones: tuple = (5, 10),
         batch_size: int = 32,
         lr: float = 1e-2,
         lr_scheduler_gamma: float = 1e-1,
         num_workers: int = 6,
-        **kwargs,
+        **_,
     ) -> None:
         """
         Args:
             dl_path: Path where the data will be downloaded
         """
         super().__init__()
         self.backbone = backbone
-        self.train_bn = train_bn
         self.milestones = milestones
         self.batch_size = batch_size
         self.lr = lr
@@ -334,13 +332,14 @@ def main(args: argparse.Namespace) -> None:
         dl_path=os.path.join(args.root_data_path, 'data'), batch_size=args.batch_size, num_workers=args.num_workers
     )
     model = TransferLearningModel(**vars(args))
-    finetuning_callback = MilestonesFinetuning(milestones=args.milestones)
+    finetuning_callback = MilestonesFinetuning(milestones=args.milestones, train_bn=args.train_bn)
 
     trainer = pl.Trainer(
         weights_summary=None,
         progress_bar_refresh_rate=1,
         num_sanity_val_steps=0,
         tpu_cores=args.tpu_cores,
+        fast_dev_run=True,
         gpus=args.gpus,
         max_epochs=args.nb_epochs,
         callbacks=[finetuning_callback]