batched-bench : add tensor_split param

ggerganov · Jan 5, 2024 · d4fca23 · d4fca23
1 parent b684502
commit d4fca23
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/examples/batched-bench/batched-bench.cpp b/examples/batched-bench/batched-bench.cpp
@@ -88,7 +88,10 @@ int main(int argc, char ** argv) {
 
     llama_model_params model_params = llama_model_default_params();
 
+    const std::vector<float> t_split (LLAMA_MAX_DEVICES, 0.0f);
+
     model_params.n_gpu_layers = n_gpu_layers;
+    model_params.tensor_split = t_split.data();
 
     llama_model * model = llama_load_model_from_file(params.model.c_str(), model_params);