Update uops.py

SzymonOzog · Mar 14, 2024 · 0aed6d3 · 0aed6d3
1 parent 55b641d
commit 0aed6d3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py
@@ -162,7 +162,7 @@ def remove_childless(self, keep:Set[UOp]):
         for vu in ru.vin:
           has_child.add(vu)
       nu: List[UOp] = [x for x in self.uops if x in has_child or x in keep]
-      self.saved_exprs = {k:v for k,v in self.saved_exprs.items() if k in nu}
+      self.saved_exprs = {k:v for k,v in self.saved_exprs.items() if v in nu}
       if len(nu) == len(self.uops): break
       if DEBUG >= 4: print(f"reduced UOp count from {len(self.uops)} to {len(nu)}")
       self.uops = nu