diff --git a/bench.cpp b/bench.cpp
index 5d1da8f..5a0e2d7 100644
--- a/bench.cpp
+++ b/bench.cpp
@@ -130,7 +130,7 @@ int main(int argc, char* argv[]) {
     std::cout << "Running benchmark..." << std::endl;
 
     // 2. Perform searches
-    AStarSearch<MapSearchNode> astarsearch;
+    AStarSearch<MapSearchNode> astarsearch((MAP_WIDTH * MAP_HEIGHT) + 1);
 
     unsigned int successes = 0;
     unsigned int failures = 0;
diff --git a/bench/run_benchmarks.py b/bench/run_benchmarks.py
index a68925e..06f6525 100755
--- a/bench/run_benchmarks.py
+++ b/bench/run_benchmarks.py
@@ -43,8 +43,8 @@ def main():
             # 1. Checkout commit
             run_cmd(f"git checkout {commit_hash}")
 
-            # 2. Inject modern bench.cpp
-            run_cmd("cp /tmp/new_bench.cpp bench.cpp")
+            # 2. Inject modern bench.cpp (from optimizations-2)
+            run_cmd(f"git checkout {original_branch} -- bench.cpp")
 
             # 3. Build release binary
             print("Building bench target...")