From 811beb0475df42f375a700e910dce06794f791b7 Mon Sep 17 00:00:00 2001 From: Theo Morales Date: Sat, 30 Jul 2022 17:13:52 +0100 Subject: [PATCH] Fix load_state_dict() to handle None layer LRs --- learn2learn/optim/transforms/layer_step_lr_transform.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/learn2learn/optim/transforms/layer_step_lr_transform.py b/learn2learn/optim/transforms/layer_step_lr_transform.py index b265f921..6d450e82 100644 --- a/learn2learn/optim/transforms/layer_step_lr_transform.py +++ b/learn2learn/optim/transforms/layer_step_lr_transform.py @@ -101,7 +101,8 @@ def load_state_dict(self, lr_state_dicts: Dict[str, Dict[str, Any]]): type(lr_state_dicts) is dict ), "Argument lr_state_dicts must be a dictionary!" for layer_name, state_dict in lr_state_dicts.items(): - self._lslr[layer_name].load_state_dict(state_dict) + if self._lslr[layer_name] is not None: + self._lslr[layer_name].load_state_dict(state_dict) def state_dict(self) -> Dict[str, Optional[Dict[str, Any]]]: return {