diff --git a/runtime/util/search.cpp b/runtime/util/search.cpp
index 4c52f5abf..901db51aa 100644
--- a/runtime/util/search.cpp
+++ b/runtime/util/search.cpp
@@ -72,7 +72,11 @@ std::unordered_set<block *, hash_block, k_eq> take_search_steps(
   while (!states.empty() && depth != 0) {
     state = states.front();
     states.pop_front();
-    states_set.erase(state);
+    if (states.empty()) {
+      states_set.clear();
+    } else {
+      states_set.erase(state);
+    }
 
     if (depth > 0) {
       depth--;