Fix tests

Author: aobolensk

src/plugins/intel_cpu/tests/unit/snippets_transformations/x64/lowered/buffer_allocation.cpp

@@ -191,7 +191,7 @@ class MHABF16AMXBufferAllocationTest : public BufferAllocationCPUTest {
     std::shared_ptr<ov::Model> GetModel(const std::vector<ov::PartialShape>& shapes) const override {
         const auto subtensor_scalar = std::vector<size_t>{1};
         const auto subtensor_power = std::vector<size_t>{1, ov::snippets::utils::get_full_dim_value()};
-        const auto subtensor_full = std::vector<size_t>(2, ov::snippets::utils::get_full_dim_value());
+        const auto subtensor_flat = std::vector<size_t>(1, ov::snippets::utils::get_full_dim_value());
 
         OPENVINO_ASSERT(shapes.size() == 3, "Incorrect count of input shapes");
         const auto parameter0 = std::make_shared<ov::op::v0::Parameter>(ov::element::bf16, shapes[0]);
@@ -238,12 +238,12 @@ class MHABF16AMXBufferAllocationTest : public BufferAllocationCPUTest {
         MarkOp(store, subtensor_scalar);
         MarkOp(power, subtensor_power);
 
-        MarkOp(brgemm_cpu0, subtensor_full);
-        MarkOp(brgemm_cpu1, subtensor_full);
-        MarkOp(brgemm_copyb0, subtensor_full);
-        MarkOp(brgemm_copyb1, subtensor_full);
-        MarkOp(scratch0, subtensor_full);
-        MarkOp(scratch1, subtensor_full);
+        MarkOp(brgemm_cpu0, subtensor_flat);
+        MarkOp(brgemm_cpu1, subtensor_flat);
+        MarkOp(brgemm_copyb0, subtensor_flat);
+        MarkOp(brgemm_copyb1, subtensor_flat);
+        MarkOp(scratch0, subtensor_flat);
+        MarkOp(scratch1, subtensor_flat);
 
         ov::snippets::lowered::PortDescriptorUtils::get_port_descriptor_ptr(load_reshape->input(0))->set_layout(order);