Fix load_state_dict() to handle None layer LRs

learnables · Jul 30, 2022 · 811beb0 · 811beb0
1 parent c4ea713
commit 811beb0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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 {