Fix batch_size

PINTO0309 · Sep 11, 2024 · 6b5b1f4 · 6b5b1f4
1 parent ad0a96e
commit 6b5b1f4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnx2tf/ops/Flatten.py b/onnx2tf/ops/Flatten.py
@@ -91,7 +91,7 @@ def make_node(
         and len(graph_node_output.shape) == 2 \
         and axis == input_tensor_rank - 1 \
         and not isinstance(graph_node_output.shape[0], str):
-        cal_shape = (1, -1)
+        cal_shape = (graph_node_output.shape[0], -1)
     elif graph_node_output.shape is not None \
         and len(graph_node_output.shape) == 2 \
         and axis == input_tensor_rank - 1 \