From cccf7de0ae1c119231b9aade5ca4184fc80e70b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Mon, 8 Mar 2021 02:58:03 +0100
Subject: [PATCH] Fix trainer not resetting lightning_optimizers (#6372)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>
---
 CHANGELOG.md                            |  3 +++
 pytorch_lightning/trainer/optimizers.py |  5 ++++-
 tests/trainer/test_trainer.py           | 25 +++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9ba5699913960..e759e89235cf8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Ensure that clip gradients is only called if the value is greater than 0 ([#6330](https://github.com/PyTorchLightning/pytorch-lightning/pull/6330)
 
 
+- Fixed `Trainer` not resetting `lightning_optimizers` when calling `Trainer.fit()` multiple times ([#6372](https://github.com/PyTorchLightning/pytorch-lightning/pull/6372))
+
+
 ## [1.2.2] - 2021-03-02
 
 ### Added
diff --git a/pytorch_lightning/trainer/optimizers.py b/pytorch_lightning/trainer/optimizers.py
index ea881b796e825..a247fb92cd22f 100644
--- a/pytorch_lightning/trainer/optimizers.py
+++ b/pytorch_lightning/trainer/optimizers.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC
-from typing import List, Optional, Tuple, Dict, Any
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 from torch import optim
@@ -27,7 +27,10 @@
 
 class TrainerOptimizersMixin(ABC):
 
+    _lightning_optimizers: Optional[List[LightningOptimizer]]
+
     def init_optimizers(self, model: LightningModule) -> Tuple[List, List, List]:
+        self._lightning_optimizers = None
         optim_conf = model.configure_optimizers()
         if optim_conf is None:
             rank_zero_warn(
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 24a6a90fc620d..59f3c2b54c13c 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -1825,3 +1825,28 @@ def backward(self, *args, **kwargs):
         "training_step",
         "backward",
     ]
+
+
+def test_init_optimizers_resets_lightning_optimizers(tmpdir):
+    """ Test that the Trainer resets the `lightning_optimizers` list everytime new optimizers get initialized. """
+
+    def compare_optimizers():
+        assert trainer.lightning_optimizers[0].optimizer is trainer.optimizers[0]
+
+    model = BoringModel()
+    model.lr = 0.2
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        auto_lr_find=True,
+    )
+
+    trainer.tune(model)
+    compare_optimizers()
+
+    trainer.fit(model)
+    compare_optimizers()
+
+    trainer.max_epochs = 2  # simulate multiple fit calls
+    trainer.fit(model)
+    compare_optimizers()