fix test_onnx.py error

Change-Id: I1f8798a9c0533db545985595979fecbbde2dda25
sophgo · Feb 15, 2023 · f26eb4b · f26eb4b
1 parent 57f4a87
commit f26eb4b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/python/test/test_onnx.py b/python/test/test_onnx.py
@@ -24,7 +24,7 @@
 import multiprocessing
 
 BM1684X_Failed_Cases = [
-    "PadAvgPool2d", "PadEdge", "QDQ", "QDQConv", "TopK", "TorchArgMax", "TorchActivation",
+    "PadAvgPool2d", "PadEdge", "QDQ", "QDQConv", "TopK", "TorchArgmax", "TorchActivation",
     "TorchChannelShuffle"
 ]
 CV18XX_Failed_Cases = [

diff --git a/python/utils/mlir_parser.py b/python/utils/mlir_parser.py
@@ -111,8 +111,12 @@ def loc(op):
 
     @staticmethod
     def shape(op):
-        shape_type = mlir.ir.ShapedType(op.results[0].type)
-        shape = [shape_type.get_dim_size(i) for i in range(shape_type.rank)]
+        shape = []
+        for result in op.results:
+            if str(result.type) != 'none':
+                shape_type = mlir.ir.ShapedType(result.type)
+                shape = [shape_type.get_dim_size(i) for i in range(shape_type.rank)]
+                break
         return shape
 
     @staticmethod