From aedd4685fa07fb6a20d6d1d754b5a01ba9876ead Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 18 Apr 2023 17:36:31 -0700 Subject: [PATCH] huh...this is faster --- tinygrad/nn/optim.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index 0fd94723a1..0a2d73e63a 100644 --- a/tinygrad/nn/optim.py +++ b/tinygrad/nn/optim.py @@ -30,9 +30,9 @@ class SGD(Optimizer): def step(self) -> None: for i, t in enumerate(self.params): assert t.grad is not None - g = t.grad + g = t.grad.realize() if self.momentum: - self.b[i].assign(self.momentum * self.b[i] + g).realize() # NOTE: self.b[i] is zero on the first run, no if required + self.b[i].assign((self.momentum * self.b[i] + g).realize()) # NOTE: self.b[i] is zero on the first run, no if required g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i] t.assign((t.detach() - g * self.lr).realize()) self.realize(self.b)