From 133fd826d2b4cb2ceb8c8b8ce07d7a749f01e7f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Mon, 15 Mar 2021 19:06:17 +0100 Subject: [PATCH] fix attribute access in LightningModule.toggle_optimizer (#6513) (cherry picked from commit ea36ee30b0a2eec1da693b42739aafed0389fbb9) --- pytorch_lightning/core/lightning.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py index 00611d87d7f35..6a83b7b1f8637 100644 --- 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,