From 785b4ea8ac5ee3cfa8eb33bd70a3499e8ceef65d Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 23 Jun 2025 13:11:19 -0400 Subject: [PATCH] optim flatten().shape[0] is numel (#10935) --- tinygrad/nn/optim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index 19d453ffa1..d08fec3b80 100644 --- a/tinygrad/nn/optim.py +++ b/tinygrad/nn/optim.py @@ -21,7 +21,7 @@ class Optimizer: # store lr in at least float32 precision self.lr = Tensor(lr if getenv("CONST_LR") else [lr], requires_grad=False, device=self.device, dtype=least_upper_dtype(dtypes.default_float, dtypes.float32)) - if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.flatten().shape[0], initial=0)) + if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.numel(), initial=0)) def _new_optim_param(self) -> list[Tensor]: param_dtype = getenv("OPTIM_DTYPE", "float32")