diff --git a/BUILD.bazel b/BUILD.bazel
index 70f96bc..bbcaa2f 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -167,6 +167,6 @@ cc_binary(
         "@hwy//:nanobenchmark",
         "@hwy//:profiler",
         "@hwy//:thread_pool",
-        "//third_party/json",
+        "@nlohmann_json//:json",
     ],
 )
diff --git a/benchmark.cc b/benchmark.cc
index db44e97..d71bbc6 100644
--- a/benchmark.cc
+++ b/benchmark.cc
@@ -3,18 +3,20 @@
 #include <iostream>
 #include <ostream>
 #include <sstream>
+#include <string>
+#include <utility>  // std::pair
 #include <vector>
 
 #include "nlohmann/json.hpp"
 // copybara:import_next_line:gemma_cpp
+#include "gemma.h"
+// copybara:import_next_line:gemma_cpp
 #include "util/app.h"
 // copybara:import_next_line:gemma_cpp
 #include "util/args.h"
 #include "hwy/base.h"
 #include "hwy/contrib/thread_pool/thread_pool.h"
 #include "hwy/highway.h"
-#include "hwy/per_target.h"
-#include "hwy/profiler.h"
 #include "hwy/timer.h"
 
 using json = nlohmann::json;