Merge pull request chainer#8332 from msakai/fix-onnx-chainer-separate…

…-output Fix onnx_chainer's exporter of Separate to handle single output case
niboshi · Oct 29, 2019 · 273b533 · 273b533
2 parents 1146a0b + 86c66f2
commit 273b533
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/onnx_chainer/functions/array.py b/onnx_chainer/functions/array.py
@@ -509,6 +509,8 @@ def convert_Separate(func, opset_version, input_names, output_names, context):
     gb = onnx_helper.GraphBuilder()
     split_outs = gb.op(
         'Split', input_names, num_outputs=len(output_names), axis=func.axis)
+    if len(output_names) == 1:
+        split_outs = [split_outs]
     for i, node_name in enumerate(split_outs):
         gb.op_output_named(
             'Squeeze', [node_name], [output_names[i]], axes=[func.axis])

diff --git a/tests/onnx_chainer_tests/functions_tests/test_arrays.py b/tests/onnx_chainer_tests/functions_tests/test_arrays.py
@@ -180,6 +180,8 @@
      'input_argname': 'x', 'args': {}, 'name': 'separate_axis0'},
     {'ops': 'separate', 'input_shape': (2, 3),
      'input_argname': 'x', 'args': {'axis': 1}, 'name': 'separate_axis1'},
+    {'ops': 'separate', 'input_shape': (1, 2, 3),
+     'input_argname': 'x', 'args': {}, 'name': 'separate_single_output'},
 
     # moveaxis
     {'ops': 'moveaxis', 'input_shape': (2, 3, 4, 5),