From f26eb4beb373c157546777c0031693bfa8b92595 Mon Sep 17 00:00:00 2001
From: "pengchao.hu" <pengchao.hu@sophgo.com>
Date: Wed, 15 Feb 2023 12:21:51 +0800
Subject: [PATCH] fix test_onnx.py error

Change-Id: I1f8798a9c0533db545985595979fecbbde2dda25
---
 python/test/test_onnx.py    | 2 +-
 python/utils/mlir_parser.py | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/test/test_onnx.py b/python/test/test_onnx.py
index 68abe5fed..19b9d5f1f 100755
--- 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
index 063bd3f40..d77b55e43 100755
--- 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