simplify teardown with finally block

Lightning-AI · Jun 14, 2020 · f744285 · f744285
1 parent d99bd22
commit f744285
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
@@ -359,7 +359,6 @@ def train(self):
                 self.run_training_epoch()
 
                 if self.max_steps and self.max_steps == self.global_step:
-                    self.run_training_teardown()
                     return
 
                 # update LR schedulers
@@ -380,11 +379,8 @@ def train(self):
                         # stop training
                         stop = should_stop and met_min_epochs
                         if stop:
-                            self.run_training_teardown()
                             return
 
-            self.run_training_teardown()
-
         except KeyboardInterrupt:
             rank_zero_warn('Detected KeyboardInterrupt, attempting graceful shutdown...')
 
@@ -394,8 +390,8 @@ def train(self):
 
                 for proc in self.interactive_ddp_procs:
                     subprocess.Popen.kill(proc)
-
-                self.run_training_teardown()
+        finally:
+            self.run_training_teardown()
 
     def run_training_epoch(self):