comment

neuralmagic · Feb 5, 2024 · 6137ef2 · 6137ef2
1 parent d353357
commit 6137ef2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/sparsezoo/utils/onnx/analysis.py b/src/sparsezoo/utils/onnx/analysis.py
@@ -392,9 +392,9 @@ def get_node_weight(
         raise Exception(f"Parameter for {node.name} not found")
 
     # some weights are not accessible, and returns the zero_points, which are scalars
-    if weight is not None and weight.size > 1:
-        return weight
-    return None
+    if weight is not None and weight.ndim == 0:
+        return None
+    return weight
 
 
 def get_layer_and_op_counts(