bugfixes

PINTO0309 · Oct 4, 2022 · a0cb1ad · a0cb1ad
1 parent 5bb8754
commit a0cb1ad
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/onnx2tf/ops/Unsqueeze.py b/onnx2tf/ops/Unsqueeze.py
@@ -65,14 +65,14 @@ def make_node(
         axes = [
             convert_axis(
                 axis=idx,
-                tensor_rank=tensor_rank,
+                tensor_rank=tensor_rank+len(axes),
                 before_op_output_shape_trans=before_op_output_shape_trans,
             ) for idx in axes
         ]
     elif axes is not None and isinstance(axes, np.ndarray) and len(axes.shape) == 0:
         axes = convert_axis(
             axis=axes,
-            tensor_rank=tensor_rank,
+            tensor_rank=tensor_rank+1,
             before_op_output_shape_trans=before_op_output_shape_trans,
         )
         axes = list(axes[np.newaxis])