fix: resume optimizer LR, step display, and consumed_samples alignment

2026-07-15 08:59:52 +08:00
parent 13998da15a
commit bb175fda91
3 changed files with 8 additions and 2 deletions
@@ -90,6 +90,7 @@ class MuonMix(optim.Optimizer):
    def load_state_dict(self, state_dict: Dict[str, Any]):
        self.muon.load_state_dict(state_dict["muon"])
        self.adamw.load_state_dict(state_dict["adamw"])
+        self.param_groups = [*self.muon.param_groups, *self.adamw.param_groups]


 def parse_args() -> argparse.Namespace: