From 96d25e95f2e40cb57c28067a1e4eeab81eff7dc7 Mon Sep 17 00:00:00 2001 From: shihchengli Date: Wed, 31 Jul 2024 23:15:49 -0400 Subject: [PATCH] Loading `train_dataloader` before estimating `max_batches` --- src/lightning/pytorch/trainer/trainer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index 406f686efe732..892e5a1457f32 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/src/lightning/pytorch/trainer/trainer.py @@ -1518,6 +1518,9 @@ def predict_dataloaders(self) -> Optional[EVAL_DATALOADERS]: @property def num_training_batches(self) -> Union[int, float]: """The number of training batches that will be used during ``trainer.fit()``.""" + if self.train_dataloader is None: + rank_zero_info("Loading `train_dataloader` to estimate number of training batches.") + self.fit_loop.setup_data() return self.fit_loop.max_batches @property