Skip to content

Commit

Permalink
Add tensor_meta to getitem nodes, this was preventing generating shlo…
Browse files Browse the repository at this point in the history
… for graphs with multiple outputs
  • Loading branch information
AleksKnezevic committed Dec 6, 2024
1 parent ab4ce10 commit 6829950
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions tt_torch/dynamo/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,12 +226,12 @@ def compile_op(self, node, *inputs, **kwargs):
):
getitem_nodes = []
graph_node.meta["val"] = node.meta["val"]
for idx, _ in enumerate(node.meta["tensor_meta"]):
for idx, tensor_meta in enumerate(node.meta["tensor_meta"]):
getitem_node = graph.call_function(
operator.getitem, args=(graph_node, idx)
)
# getitem_node.meta["val"] = graph_node.meta["val"]
getitem_nodes.append(getitem_node)
getitem_node.meta["tensor_meta"] = tensor_meta
out = graph.output(tuple(getitem_nodes))
else:
out = graph.output((graph_node,))
Expand Down

0 comments on commit 6829950

Please sign in to comment.