Motif-Technologies
/

optimizer

Model card Files Files and versions

fix init

#7

by leejunhyeok - opened Sep 10, 2025

base: refs/heads/main

←

from: refs/pr/7

Discussion Files changed

Files changed (1) hide show

torch-ext/optimizer/muon.py +2 -2

torch-ext/optimizer/muon.py CHANGED Viewed

@@ -210,7 +210,7 @@ class Muon(torch.optim.Optimizer):
     def __init__(
         self,
-        model,
         is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
@@ -233,7 +233,7 @@ class Muon(torch.optim.Optimizer):
             none_grad=none_grad,
         )
-        super().__init__(model.parameters(), defaults)
         self.is_muon_func = is_muon_func
         self.model = model

     def __init__(
         self,
+        params,
         is_muon_func=default_is_muon,
         lr=1e-3,
         momentum=0.95,
             none_grad=none_grad,
         )
+        super().__init__(params, defaults)
         self.is_muon_func = is_muon_func
         self.model = model