Merge pull request #16 from DaoyuanLi2816/fix-scheduler-state-dict-optimizer-key

Fix scheduler state_dict to drop the optimizer key by equality
2026-06-30 14:42:21 +08:00
parent af59e17a37 a71367a787
commit afdfa7c938
1 changed files with 1 additions and 1 deletions
@@ -11,7 +11,7 @@ class TwoStageScheduler(_LRScheduler):

    def state_dict(self):
        state_dict = {
-            key: value for key, value in self.__dict__.items() if key not in "optimizer"
+            key: value for key, value in self.__dict__.items() if key != "optimizer"
        }
        if isinstance(state_dict["after_scheduler"], _LRScheduler):
            state_dict["after_scheduler_type"] = type(