Fix MO multinomial shape array (openvinotoolkit#20548)

praasz · Oct 19, 2023 · dc9ec1e · dc9ec1e
1 parent 46935e0
commit dc9ec1e
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tools/mo/openvino/tools/mo/ops/multinomial.py b/tools/mo/openvino/tools/mo/ops/multinomial.py
@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from openvino.tools.mo.front.common.partial_infer.utils import dynamic_dimension
+from openvino.tools.mo.front.common.partial_infer.utils import dynamic_dimension_value, shape_array
 from openvino.tools.mo.front.extractor import bool_to_str
 from openvino.tools.mo.graph.graph import Graph, Node
 
@@ -63,7 +63,7 @@ def infer(node: Node):
 
         num_samples = node.in_port(1).data.get_value()
         if num_samples is not None:
-            output_shape.append(num_samples)
+            output_shape.append(np.array(num_samples).item())
         else:
-            output_shape.append(dynamic_dimension)
-        node.out_port(0).data.set_shape(output_shape)
+            output_shape.append(dynamic_dimension_value)
+        node.out_port(0).data.set_shape(shape_array(output_shape))