Set axes=[0] for apply_squeeze rather than leaving unset (#668)

onnx · Dec 17, 2020 · 16a42f5 · 16a42f5
1 parent 6d799ad
commit 16a42f5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/keras2onnx/ke2onnx/lstm.py b/keras2onnx/ke2onnx/lstm.py
@@ -224,7 +224,7 @@ def build_output_states(scope, operator, container, output_names, bidirectional=
             squeeze_names.extend(list(zip(split_names, outputs)))
 
         for split_name, output_name in squeeze_names:
-            apply_squeeze(scope, split_name, output_name, container)
+            apply_squeeze(scope, split_name, output_name, container, axes=[0])
 
     else:
         output_state = op.return_state
@@ -234,8 +234,8 @@ def build_output_states(scope, operator, container, output_names, bidirectional=
 
         output_h = operator.outputs[1].full_name
         output_c = operator.outputs[2].full_name
-        apply_squeeze(scope, lstm_h, output_h, container)
-        apply_squeeze(scope, lstm_c, output_c, container)
+        apply_squeeze(scope, lstm_h, output_h, container, axes=[0])
+        apply_squeeze(scope, lstm_c, output_c, container, axes=[0])
 
 
 def _calculate_keras_lstm_output_shapes(operator):

diff --git a/keras2onnx/ke2onnx/simplernn.py b/keras2onnx/ke2onnx/simplernn.py
@@ -391,14 +391,14 @@ def build_output_states(scope, operator, container, output_names, bidirectional=
             apply_split(scope, rnn_h, split_names, container)
 
             for split_name, output_name in zip(split_names, output_names):
-                apply_squeeze(scope, split_name, output_name, container)
+                apply_squeeze(scope, split_name, output_name, container, axes=[0])
 
     else:
         output_state = op.return_state
 
         if output_state:
             output_h = operator.outputs[1].full_name
-            apply_squeeze(scope, rnn_h, output_h, container)
+            apply_squeeze(scope, rnn_h, output_h, container, axes=[0])
 
 
 def is_time_major(op, bidirectional):