From bcfcd60f5531b596ab919e95bcd0ab78c888143b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 31 Mar 2025 11:02:03 +0800 Subject: [PATCH] opt weights --- tinygrad/engine/jit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index dfd5556456..59729e127e 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -314,6 +314,7 @@ class TinyJit(Generic[ReturnType]): # set this for next run self.captured = CapturedJit(ret, jit_cache, input_replace, extra_view_inputs, names, st_vars_dtype_device) + self.captured.optimize_weights() elif self.cnt >= 2: # jit exec assert self.captured is not None