diff options
Diffstat (limited to 'ep_run/muon.py')
| -rw-r--r-- | ep_run/muon.py | 6 |
1 files changed, 5 insertions, 1 deletions
diff --git a/ep_run/muon.py b/ep_run/muon.py index 3116030..7281811 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -38,12 +38,16 @@ class Muon(torch.optim.Optimizer): class MultiOpt: - """duck-typed bundle of optimizers (step/zero_grad API-compatible).""" + """duck-typed bundle of optimizers (step/zero_grad/state_dict API-compatible).""" def __init__(self, opts): self.optimizers = opts def step(self): for o in self.optimizers: o.step() def zero_grad(self, set_to_none=True): for o in self.optimizers: o.zero_grad(set_to_none=set_to_none) + def state_dict(self): + return [o.state_dict() for o in self.optimizers] + def load_state_dict(self, sds): + for o, sd in zip(self.optimizers, sds): o.load_state_dict(sd) class MultiSched: |