Remove Strategy.optimizer_zero_grad (#11246)

Lightning-AI · Jan 3, 2022 · b082715 · b082715
1 parent 4eede7c
commit b082715
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -346,6 +346,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Removed `Strategy.init_optimizers` in favor of `Strategy.setup_optimizers` ([#11236](https://github.com/PyTorchLightning/pytorch-lightning/pull/11236))
 
+
+- Removed `Strategy.optimizer_zero_grad` ([#11246](https://github.com/PyTorchLightning/pytorch-lightning/pull/11246))
+
 ### Fixed
 
 - Fixed security vulnerabilities CVE-2020-1747 and CVE-2020-14343 caused by the `PyYAML` dependency ([#11099](https://github.com/PyTorchLightning/pytorch-lightning/pull/11099))

diff --git a/pytorch_lightning/loops/optimization/optimizer_loop.py b/pytorch_lightning/loops/optimization/optimizer_loop.py
@@ -400,7 +400,7 @@ def _optimizer_zero_grad(self, batch_idx: int, optimizer: torch.optim.Optimizer,
             optimizer: the current optimizer
             opt_idx: the index of the current optimizer
         """
-        self.trainer._call_strategy_hook(
+        self.trainer._call_lightning_module_hook(
             "optimizer_zero_grad", self.trainer.current_epoch, batch_idx, optimizer, opt_idx
         )
         self.optim_progress.optimizer.zero_grad.increment_completed()

diff --git a/pytorch_lightning/strategies/strategy.py b/pytorch_lightning/strategies/strategy.py
@@ -177,10 +177,6 @@ def optimizer_step(
         model = model or self.lightning_module
         self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
 
-    def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Optimizer, opt_idx: int) -> None:
-        """Zeros all model parameter's gradients."""
-        self.lightning_module.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx)
-
     def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]) -> Tuple[Module, List[Optimizer]]:
         """Setup a model and multiple optimizers together.