update

Author: deependujha

tests/tests_pytorch/models/test_hooks.py

@@ -477,6 +477,8 @@ def training_step(self, batch, batch_idx):
         # DeepSpeed needs the batch size to figure out throughput logging
         *([{"name": "train_dataloader"}] if using_deepspeed else []),
         {"name": "configure_model"},
+        {"name": "Callback.on_before_optimizer_setup", "args": (trainer, model)},
+        {"name": "on_before_optimizer_setup"},
         {"name": "configure_optimizers"},
         {"name": "Callback.on_fit_start", "args": (trainer, model)},
         {"name": "on_fit_start"},
@@ -574,6 +576,8 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_epochs(tmp_path):
         {"name": "Callback.setup", "args": (trainer, model), "kwargs": {"stage": "fit"}},
         {"name": "setup", "kwargs": {"stage": "fit"}},
         {"name": "configure_model"},
+        {"name": "Callback.on_before_optimizer_setup", "args": (trainer, model)},
+        {"name": "on_before_optimizer_setup"},
         {"name": "on_load_checkpoint", "args": (loaded_ckpt,)},
         {"name": "Callback.on_load_checkpoint", "args": (trainer, model, loaded_ckpt)},
         {"name": "Callback.load_state_dict", "args": ({"foo": True},)},
@@ -654,6 +658,8 @@ def test_trainer_model_hook_system_fit_no_val_and_resume_max_steps(tmp_path):
         {"name": "Callback.setup", "args": (trainer, model), "kwargs": {"stage": "fit"}},
         {"name": "setup", "kwargs": {"stage": "fit"}},
         {"name": "configure_model"},
+        {"name": "Callback.on_before_optimizer_setup", "args": (trainer, model)},
+        {"name": "on_before_optimizer_setup"},
         {"name": "on_load_checkpoint", "args": (loaded_ckpt,)},
         {"name": "Callback.on_load_checkpoint", "args": (trainer, model, loaded_ckpt)},
         {"name": "Callback.load_state_dict", "args": ({"foo": True},)},