allow val dataloader

Lightning-AI · May 4, 2020 · 24817e9 · 24817e9
1 parent b6c271c
commit 24817e9
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/pytorch_lightning/trainer/lr_finder.py b/pytorch_lightning/trainer/lr_finder.py
@@ -53,6 +53,7 @@ def _run_lr_finder_internally(self, model: LightningModule):
     def lr_find(self,
                 model: LightningModule,
                 train_dataloader: Optional[DataLoader] = None,
+                val_dataloaders: Optional[DataLoader] = None,
                 min_lr: float = 1e-8,
                 max_lr: float = 1,
                 num_training: int = 100,
@@ -133,7 +134,7 @@ def lr_find(self,
         self.checkpoint_callback = False
         self.early_stop_callback = None
         self.enable_early_stop = False
-        
+
         # Required for saving the model
         self.optimizers, self.schedulers = [], [],
         self.model = model
@@ -152,7 +153,9 @@ def lr_find(self,
         model.configure_optimizers = lr_finder._get_new_optimizer(optimizers[0])
 
         # Fit, lr & loss logged in callback
-        self.fit(model, train_dataloader=train_dataloader)
+        self.fit(model,
+                 train_dataloader=train_dataloader,
+                 val_dataloaders=val_dataloaders)
 
         # Prompt if we stopped early
         if self.global_step != num_training: