Fix llama_triton_example (InternLM#1414)

zyearw1024 · Apr 9, 2024 · b5b2b8b · b5b2b8b
1 parent edca3d3
commit b5b2b8b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/cpp/llama/llama_triton_example.cc b/examples/cpp/llama/llama_triton_example.cc
@@ -164,7 +164,7 @@ broadCastRequest(const std::vector<int>& v_start_ids,
                                 request_output_len_ptr}},
                 {"bad_words_list",
                  triton::Tensor{
-                     triton::MEMORY_GPU, triton::TYPE_INT32, {2, v_input_bad_words.size() / 2}, d_input_bad_words}},
+                     triton::MEMORY_GPU, triton::TYPE_INT32, {1, 2, v_input_bad_words.size() / 2}, d_input_bad_words}},
                 {"start_id",
                  triton::Tensor{triton::MEMORY_CPU, triton::TYPE_INT32, {(size_t)request_batch_size}, start_ids_ptr}},
                 {"end_id",