Fix setup callback hook to pass LightningModule through (#4608)

* Fix setup callback hook

* Update CHANGELOG.md

* Update test_trainer.py

* Update test_trainer.py

* Update test_trainer.py

* fix chlog

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

CHANGELOG.md

@@ -48,6 +48,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `setup` callback hook to correctly pass the LightningModule through ([#4608](https://github.com/PyTorchLightning/pytorch-lightning/pull/4608))
+
 
 
 

pytorch_lightning/trainer/callback_hook.py

@@ -26,10 +26,10 @@ class TrainerCallbackHookMixin(ABC):
     callbacks: List[Callback] = []
     get_model: Callable
 
-    def setup(self, stage: str):
+    def setup(self, model, stage: str):
         """Called in the beginning of fit and test"""
         for callback in self.callbacks:
-            callback.setup(self, self.get_model(), stage)
+            callback.setup(self, model, stage)
 
     def teardown(self, stage: str):
         """Called at the end of fit and test"""

pytorch_lightning/trainer/trainer.py

@@ -850,7 +850,7 @@ def call_setup_hook(self, model):
             called = self.datamodule.has_setup_test if self.testing else self.datamodule.has_setup_fit
             if not called:
                 self.datamodule.setup(stage_name)
-        self.setup(stage_name)
+        self.setup(model, stage_name)
         model.setup(stage_name)
 
     def _reset_result_and_set_hook_fx_name(self, hook_name):

tests/trainer/test_trainer.py

@@ -1431,7 +1431,8 @@ def setup(self, stage):
             self.stage = stage
 
     class TrainerSubclass(Trainer):
-        def setup(self, stage):
+        def setup(self, model, stage):
+            assert model is not None
             self.stage = stage
 
     model = CurrentModel()