Estimate stepping batches with max_steps if max_epochs is not set (#14317)

Co-authored-by: Roberto Estevão <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Adrian Wälchli <[email protected]>

Author: 4 people

src/pytorch_lightning/CHANGELOG.md

@@ -29,6 +29,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue to avoid the impact of sanity check on `reload_dataloaders_every_n_epochs` for validation ([#13964](https://github.com/Lightning-AI/lightning/pull/13964))
 
 
+- Fixed `Trainer.estimated_stepping_batches` when maximum number of epochs is not set ([#14317](https://github.com/Lightning-AI/lightning/pull/14317))
+
+
 ## [1.7.2] - 2022-08-17
 
 ### Added

src/pytorch_lightning/trainer/trainer.py

@@ -2769,8 +2769,8 @@ def configure_optimizers(self):
             )
 
         # infinite training
-        if self.max_epochs == -1 and self.max_steps == -1:
-            return float("inf")
+        if self.max_epochs == -1:
+            return float("inf") if self.max_steps == -1 else self.max_steps
 
         if self.train_dataloader is None:
             rank_zero_info("Loading `train_dataloader` to estimate number of stepping batches.")

tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py

@@ -95,9 +95,9 @@ def test_num_stepping_batches_infinite_training():
     assert trainer.estimated_stepping_batches == float("inf")
 
 
-def test_num_stepping_batches_with_max_steps():
+@pytest.mark.parametrize("max_steps", [2, 100])
+def test_num_stepping_batches_with_max_steps(max_steps):
     """Test stepping batches with `max_steps`."""
-    max_steps = 2
     trainer = Trainer(max_steps=max_steps)
     model = BoringModel()
     trainer.fit(model)