diff --git a/examples/sweep-bench/sweep-bench.cpp b/examples/sweep-bench/sweep-bench.cpp
index 865f65f0..ef3b3d59 100644
--- a/examples/sweep-bench/sweep-bench.cpp
+++ b/examples/sweep-bench/sweep-bench.cpp
@@ -94,7 +94,7 @@ int main(int argc, char ** argv) {
     };
 
     const unsigned int pp = params.n_ubatch;
-    const unsigned int tg = params.n_ubatch / 4;
+    const unsigned int tg = params.n_predict > 0 ? params.n_predict : params.n_ubatch / 4;
 
     if (!params.sweep_bench_output_jsonl) {
         LOG_TEE("\n");