fix attribute access in LightningModule.toggle_optimizer (#6513)

(cherry picked from commit ea36ee3)
Lightning-AI · Mar 16, 2021 · c8fb646 · c8fb646
1 parent 342d097
commit c8fb646
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -105,6 +105,7 @@ def __init__(self, *args, **kwargs):
         self._current_dataloader_idx = None
         self.running_stage = None
         self._automatic_optimization: bool = True
+        self._param_requires_grad_state = dict()
 
     def optimizers(self, use_pl_optimizer: bool = True) -> Union[Optimizer, List[Optimizer], List[LightningOptimizer]]:
         if use_pl_optimizer:
@@ -1295,7 +1296,7 @@ def untoggle_optimizer(self, optimizer_idx: int):
                         if param in self._param_requires_grad_state:
                             param.requires_grad = self._param_requires_grad_state[param]
         # save memory
-        del self._param_requires_grad_state
+        self._param_requires_grad_state = dict()
 
     def optimizer_step(
         self,