diff --git a/egg-herbie/src/lib.rs b/egg-herbie/src/lib.rs
index 27e96c75d..95f357121 100644
--- a/egg-herbie/src/lib.rs
+++ b/egg-herbie/src/lib.rs
@@ -203,7 +203,7 @@ pub unsafe extern "C" fn egraph_run(
     // E.g. if an e-class represents the number 2, it's safe to keep the `(2)` e-node but remove the
     // `(+ 1 1)` e-node.
     context.runner.egraph.classes_mut().for_each(|eclass| {
-        if eclass.data.is_some() {
+        if eclass.nodes.iter().any(|n| n.is_leaf()) {
             eclass.nodes.retain(|n| n.is_leaf());
         }
     });