hanzoai/deepspec
2
0
Fork 0
mirror of https://github.com/hanzoai/deepspec.git synced 2026-09-30 21:41:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix scheduler state_dict to drop the optimizer key by equality

Browse Source 
TwoStageScheduler.state_dict() used `key not in "optimizer"`, which tests
whether each key is a substring of the string "optimizer" rather than
comparing it against "optimizer". It only drops the intended "optimizer"
key today because no other scheduler attribute name happens to be a
substring of "optimizer".

Use `key != "optimizer"`, matching torch's LRScheduler.state_dict idiom and
the SpecForge optimizer code this class is adapted from.
This commit is contained in:
Daoyuan Li
2026-06-29 21:23:02 -07:00
parent ba73adaa6d
commit 1c1d993895
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deepspec/utils/optim.py
+1 -1
View File
@@ -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(