#8520: Add optional output tensors to create head decode

tenstorrent · May 15, 2024 · 42e0664 · 42e0664
1 parent b6ad9e8
commit 42e0664
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tt_eager/tt_dnn/op_library/nlp_tms/nlp_tms.hpp b/tt_eager/tt_dnn/op_library/nlp_tms/nlp_tms.hpp
@@ -151,7 +151,7 @@ inline std::vector<Tensor> nlp_create_qkv_heads_decode(
                                         Tensor(operation::get_workers_for_op_output({input_tensor})),
                                         Tensor(operation::get_workers_for_op_output({input_tensor}))};
     operation::launch_op(
-        [num_heads, num_kv_heads, mem_config] (std::vector<Tensor> input_tensors, const std::vector<std::optional<const Tensor>>& optional_input_tensors) mutable -> std::vector<Tensor> {
+        [num_heads, num_kv_heads, mem_config] (std::vector<Tensor> input_tensors, const std::vector<std::optional<const Tensor>>& optional_input_tensors, const std::vector<std::optional<Tensor>>& optional_output_tensors) mutable -> std::vector<Tensor> {
             const uint32_t num_kv_heads_val = num_kv_heads.value_or(num_heads);
 
             // Infer head_dim