nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

use float32 for sgd momentum (#10387)

Browse Source
This commit is contained in:
George Hotz
2025-05-17 21:56:44 -07:00
committed by GitHub
parent 305a3231c4
commit c91f2c4580
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/nn/optim.py
+2 -1
View File
@@ -75,7 +75,8 @@ class LARS(Optimizer):
def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001):
super().__init__(params, lr)
self.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, tcoef
self.b = [Tensor.zeros(*t.shape, dtype=t.dtype, device=t.device, requires_grad=False) for t in self.params] if self.momentum else []
self.b = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params] \
if self.momentum else []
def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]:
for i, (t, g) in enumerate(zip(self.params, grads)):