diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index b3a565161cb4..43452bc152c5 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -191,7 +191,7 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]) -> Tuple[Defaul # find all reduces, and pair them to a elementwise op. if they can't be cleanly paired, force realize the reduce (or a contig child) reduce_for_op: Dict[LazyBuffer, LazyBuffer] = {} for r in allbufs: - if r != r.base or r.op not in ReduceOps or r in realizes: continue + if r.op not in ReduceOps or r in realizes: continue group: Set[LazyBuffer] = set() _recursive_group(r, r.st, r, children, realizes, reduce_for_op, group)