Fix build

Signed-off-by: Evgeniia Nugmanova <[email protected]>
jane-intel · Dec 24, 2024 · 7c183db · 7c183db
1 parent 658c6df
commit 7c183db
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/src/core/src/pass/sdpa_to_paged_attention.cpp b/src/core/src/pass/sdpa_to_paged_attention.cpp
@@ -149,21 +149,25 @@ bool ov::pass::SDPAToPagedAttention::run_on_model(const std::shared_ptr<ov::Mode
         }
     }
 
-    for (const std::string& param_name : {"beam_idx", "attention_mask"}) {
+    for (const auto& param_name : {"beam_idx", "attention_mask"}) {
         if (auto param = get_parameter(model, param_name)) {
-            if (param_name == "attention_mask" && param->output(0).get_target_inputs().size() == 1)
-                param->output(0).get_target_inputs().begin()->replace_source_output(input_ids_node->output(0));
+            auto target_inputs = param->output(0).get_target_inputs();
+            if (!strcmp(param_name, "attention_mask") && target_inputs.size() == 1 &&
+                ov::is_type<op::util::ShapeOfBase>(target_inputs.begin()->get_node())) {
+                target_inputs.begin()->replace_source_output(input_ids_node->output(0));
+                target_inputs = param->output(0).get_target_inputs();
+            }
             model->remove_parameter(param);
 
-            if (param->output(0).get_target_inputs().size() == 0) {
+            if (!target_inputs.empty()) {
                 std::stringstream consumers;
                 consumers << std::endl;
-                for (auto& input : param->output(0).get_target_inputs()) {
+                for (auto& input : target_inputs) {
                     consumers << *input.get_node() << std::endl;
                 }
-                OPENVINO_ASSERT(param->output(0).get_target_inputs().size() == 0,
+                OPENVINO_ASSERT(target_inputs.empty(),
                                 "PagedAttention transformation failed: couldn't remove ",
-                                param->output(0).get_target_inputs().size(),
+                                target_inputs.size(),
                                 " inputs of ",
                                 param_name,
                                 " input: ",