diff --git a/src/pytorch-lightning/models/trainer.py b/src/pytorch-lightning/models/trainer.py index 1088225f..7d6f43a8 100644 --- a/src/pytorch-lightning/models/trainer.py +++ b/src/pytorch-lightning/models/trainer.py @@ -267,7 +267,8 @@ class Trainer(TrainerIO): break # give model a chance to end epoch early - if self.model.should_stop_epoch(): + if self.model.should_stop_epoch: + self.model.should_stop_epoch = False break # --------------- diff --git a/src/pytorch-lightning/root_module/hooks.py b/src/pytorch-lightning/root_module/hooks.py index e5a62f01..0d903f66 100644 --- a/src/pytorch-lightning/root_module/hooks.py +++ b/src/pytorch-lightning/root_module/hooks.py @@ -19,5 +19,3 @@ class ModelHooks(torch.nn.Module): def on_post_performance_check(self): pass - def should_stop_epoch(self): - return False diff --git a/src/pytorch-lightning/root_module/root_module.py b/src/pytorch-lightning/root_module/root_module.py index ab49a9a1..b12a2087 100644 --- a/src/pytorch-lightning/root_module/root_module.py +++ b/src/pytorch-lightning/root_module/root_module.py @@ -24,6 +24,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks): self.overfit = hparams.overfit self.gradient_clip = hparams.gradient_clip self.num = 2 + self.should_stop_epoch = False # track if gpu was requested for checkpointing self.on_gpu = False