Fix errors

optm.py

@@ -1,16 +1,18 @@
-from tinygrad import Tensor
+from tinygrad import Tensor,nn
 import math
 
 class CosineLR:
-    def __init__(self,optm,totalSteps,minlr):
+    def __init__(self,optm,totalSteps,maxlr,minlr):
         self.optm = optm
-        self.maxlr = optm.lr
+        self.maxlr = maxlr
         self.minlr = minlr
         self.totalSteps = totalSteps
         self.steps = 0
 
     def step(self):
-        self.optm.lr = self.minlr + 0.5 * (self.maxlr - self.minlr) * (1 + math.cos((step / self.totalSteps) * math.pi))
+        lr = self.minlr + 0.5 * (self.maxlr - self.minlr) * (1 + math.cos((self.steps / self.totalSteps) * math.pi))
+        for o in self.optm:
+            o.lr = lr
         self.optm.step()
         self.steps += 1
 
@@ -18,11 +20,11 @@ class CosineLR:
         self.optm.zero_grad()
 
 
-def llmOptimizer(params,steps,minlr):
+def llmOptimizer(params,steps,maxlr,minlr):
     muon_params = [p for p in params if len(p.shape) >= 2]
     adamw_params = [p for p in params if len(p.shape) < 2]
 
-    o1 = nn.optim.Muon(muon_params, lr=hypr["starting_lr"])
-    o2 = nn.optim.AdamW(adamw_params, lr=hypr["starting_lr"])
-    optimizer = nn.optim.OptimizerGroup([o1,o2])
-    return CosineLR(optimizer,steps,minlr)
+    o1 = nn.optim.Muon(muon_params, lr=maxlr)
+    o2 = nn.optim.AdamW(adamw_params, lr=maxlr)
+    optimizer = nn.optim.OptimizerGroup(o1,o2)
+    return CosineLR(optimizer,steps,maxlr,minlr)