Minor fixes

Author: carmocca

pytorch_lightning/trainer/progress.py

@@ -148,9 +148,6 @@ def from_defaults(cls, tracker_cls: Type[ReadyCompletedTracker], **kwargs: int)
         """Utility function to easily create an instance from keyword arguments to both ``Tracker``s."""
         return cls(total=tracker_cls(**kwargs), current=tracker_cls(**kwargs))
 
-    def reset_on_epoch(self) -> None:
-        self.current.reset()
-
     def reset_on_run(self) -> None:
         self.current.reset()
 

tests/models/test_restore.py

@@ -224,31 +224,39 @@ def test_correct_step_and_epoch(tmpdir):
     model = BoringModel()
     first_max_epochs = 2
     train_batches = 2
-    trainer = Trainer(default_root_dir=tmpdir, max_epochs=first_max_epochs, limit_train_batches=train_batches)
+    trainer = Trainer(
+        default_root_dir=tmpdir, max_epochs=first_max_epochs, limit_train_batches=train_batches, limit_val_batches=0
+    )
     assert trainer.current_epoch == 0
     assert trainer.global_step == 0
 
     trainer.fit(model)
     assert trainer.current_epoch == first_max_epochs
     assert trainer.global_step == first_max_epochs * train_batches
 
-    ckpt = str(tmpdir / "model.ckpt")
-    trainer.save_checkpoint(ckpt)
+    ckpt_path = str(tmpdir / "model.ckpt")
+    trainer.save_checkpoint(ckpt_path)
+
+    ckpt = torch.load(ckpt_path)
+    assert ckpt["epoch"] == first_max_epochs
     # TODO(@carmocca): should not need `+1`
-    assert torch.load(ckpt)["global_step"] == first_max_epochs * train_batches + 1
+    assert ckpt["global_step"] == first_max_epochs * train_batches + 1
 
-    max_epochs = 4
-    trainer = Trainer(default_root_dir=tmpdir, max_epochs=max_epochs, limit_train_batches=train_batches)
+    max_epochs = first_max_epochs + 2
+    trainer = Trainer(
+        default_root_dir=tmpdir, max_epochs=max_epochs, limit_train_batches=train_batches, limit_val_batches=0
+    )
     # the ckpt state is not loaded at this point
     assert trainer.current_epoch == 0
     assert trainer.global_step == 0
 
     class TestModel(BoringModel):
         def on_pretrain_routine_end(self) -> None:
+            assert self.trainer.current_epoch == first_max_epochs
             # TODO(@carmocca): should not need `+1`
             assert self.trainer.global_step == first_max_epochs * train_batches + 1
 
-    trainer.fit(TestModel(), ckpt_path=ckpt)
+    trainer.fit(TestModel(), ckpt_path=ckpt_path)
     assert trainer.current_epoch == max_epochs
     # TODO(@carmocca): should not need `+1`
     assert trainer.global_step == max_epochs * train_batches + 1