update

Author: deependujha

docs/source-pytorch/common/hooks.rst

@@ -88,13 +88,19 @@ with the source of each hook indicated:
     │   ├── [LightningModule]
     │   ├── [LightningModule.configure_shared_model()]
     │   ├── [LightningModule.configure_model()]
+    |   |
+    │   ├── on_before_optimizer_setup()
+    │   │   ├── [Callbacks]
+    │   │   └── [LightningModule]
+    │   |
     │   ├── Strategy.restore_checkpoint_before_setup
     │   │   ├── [LightningModule.on_load_checkpoint()]
     │   │   ├── [LightningModule.load_state_dict()]
     │   │   ├── [LightningDataModule.load_state_dict()]
     │   │   ├── [Callbacks.on_load_checkpoint()]
     │   │   └── [Callbacks.load_state_dict()]
-    │   └── [Strategy]
+    |   |
+    │   └── [Strategy] (configures optimizers, lr schedulers, precision, etc.)
     │
     ├── on_fit_start()
     │   ├── [Callbacks]

tests/tests_pytorch/callbacks/test_callback_hooks.py

@@ -64,10 +64,10 @@ def test_callback_on_before_optimizer_setup(tmp_path):
     class CB(Callback):
         def setup(self, trainer, pl_module, stage=None):
             assert len(trainer.optimizers) == 0
-            assert pl_module.layer is None  # setup is called before `LightningModule.configure_model`
+            assert pl_module.layer is None  # called before `LightningModule.configure_model`
 
         def on_before_optimizer_setup(self, trainer, pl_module):
-            assert len(trainer.optimizers) == 0
+            assert len(trainer.optimizers) == 0  # `LightningModule.configure_optimizers` hasn't been called yet
             assert pl_module.layer is not None  # called after `LightningModule.configure_model`
 
         def on_fit_start(self, trainer, pl_module):