fix numpy array type check

inisis · Nov 10, 2024 · 97bacb7 · 97bacb7
1 parent c5fb05a
commit 97bacb7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/onnxslim/core/pattern/elimination/reshape.py b/onnxslim/core/pattern/elimination/reshape.py
@@ -43,10 +43,10 @@ def check_constant_mergeable(reshape_node):
                 """
                 if isinstance(reshape_node.inputs[1], gs.Constant):
                     input_shape = reshape_node.inputs[0].shape
-                    reshape_shape = reshape_node.inputs[1].values
+                    reshape_shape = reshape_node.inputs[1].values.tolist()
                     if input_shape is not None and np.any(reshape_shape == 0):
                         shape = [
-                            input_shape[i] if dim_size == 0 else dim_size for i, dim_size in enumerate(reshape_shape)
+                            input_shape[i] if dim_size == 0 else reshape_shape[i] for i, dim_size in enumerate(reshape_shape)
                         ]
                         if not all(isinstance(item, int) for item in shape):
                             return False