Fix SDPA tests

maxnick · Apr 26, 2024 · cc2383e · cc2383e
1 parent 672a10d
commit cc2383e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/...lugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_multiple_query_sdp.cpp b/...lugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_multiple_query_sdp.cpp
@@ -269,7 +269,12 @@ class ConcatMultiQuerySDPTest : public testing::WithParamInterface<ConcatMultiQu
             outputs.push_back(copy);
         }
         auto states = inferRequest.query_state();
-        for (auto&& state : states) {
+        for (std::string name : {"pastk", "pastv"}) {
+            auto itr = std::find_if(states.begin(), states.end(), [&](const ov::VariableState& state) {
+                return name == state.get_name();
+            });
+            OPENVINO_ASSERT(itr != states.end(), "Failed to find ", name, " state");
+            const auto& state = *itr;
             auto state_tensor = state.get_state();
             ov::Tensor copy{state_tensor.get_element_type(), state_tensor.get_shape()};
             state_tensor.copy_to(copy);

diff --git a/...s/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_transpose_sdp_transpose.cpp b/...s/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_transpose_sdp_transpose.cpp
@@ -253,7 +253,12 @@ class ConcatSDPTransposeTest : public ConcatSDPTransposeTestBase {
             outputs.push_back(copy);
         }
         auto states = inferRequest.query_state();
-        for (auto&& state : states) {
+        for (std::string name : {"pastk", "pastv"}) {
+            auto itr = std::find_if(states.begin(), states.end(), [&](const ov::VariableState& state) {
+                return name == state.get_name();
+            });
+            OPENVINO_ASSERT(itr != states.end(), "Failed to find ", name, " state");
+            const auto& state = *itr;
             auto state_tensor = state.get_state();
             ov::Tensor copy{state_tensor.get_element_type(), state_tensor.get_shape()};
             state_tensor.copy_to(copy);