Conditionally add dequant axis

tenstorrent · Aug 1, 2024 · 78e10bb · 78e10bb
1 parent a14c698
commit 78e10bb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/tvm/contrib/pybuda_compile.py b/python/tvm/contrib/pybuda_compile.py
@@ -489,12 +489,13 @@ def duplicate_dequantize_nodes_in_onnx_graph(onnx_module):
             for i, consumer_name in enumerate(consumers):
                 new_node_name = node.name + f"_clone{i}"
                 new_output_name = output_name + f"_clone{i}"
+                attrs = {"axis": node.attribute[0].i} if len(node.attribute) > 0 else {}
                 cloned_node = onnx.helper.make_node(
                     node.op_type,
                     node.input,
                     [new_output_name],
                     name=new_node_name,
-                    axis=node.attribute[0].i
+                    **attrs
                 )
 
                 # Add the cloned node to the list of nodes to add