Apply review comments

openvinotoolkit · Jan 16, 2025 · afcc6fe · afcc6fe
1 parent 8ae608d
commit afcc6fe
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/core/src/op/if.cpp b/src/core/src/op/if.cpp
@@ -34,12 +34,12 @@ static ov::PartialShape resolve_shape(const ov::PartialShape& then_pshape, const
     }
     if (then_rank.get_length() != else_rank.get_length()) {
         auto is_one_element = [](const ov::PartialShape& pshape) {
-            return pshape.size() == 0 || pshape[0].get_max_length() == 1;
+            return pshape.size() == 0 || (pshape.is_static() && pshape[0].get_length() == 1);
         };
         // Union of scalar and 1D case
         if (then_rank.get_length() <= 1 && else_rank.get_length() <= 1) {
-            return (is_one_element(then_pshape) && is_one_element(else_pshape)) ?
-                    ov::PartialShape{1} : ov::PartialShape::dynamic(1);
+            return (is_one_element(then_pshape) && is_one_element(else_pshape)) ? ov::PartialShape{1}
+                                                                                : ov::PartialShape::dynamic(1);
         } else {
             return ov::PartialShape::dynamic();
         }

diff --git a/src/core/tests/type_prop/if.cpp b/src/core/tests/type_prop/if.cpp
@@ -365,7 +365,7 @@ TEST(type_prop, if_output_one_element) {
     if_op->set_input(Y, nullptr, Ye);
     auto res = if_op->set_output(then_body_res, else_body_res);
     auto result0 = make_shared<ov::op::v0::Result>(res);
-    PartialShape out_shape{PartialShape{1}};
+    PartialShape out_shape{1};
     auto sh = result0->get_output_partial_shape(0);
     EXPECT_EQ(sh, out_shape);
 }