Skip to content

Commit 9beac22

Browse files
authored
chore: fix attn test (#3055)
1 parent 411ad4b commit 9beac22

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

py/torch_tensorrt/dynamo/lowering/passes/lower_scaled_dot_product_attention.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,10 @@ def lower_scaled_dot_product_attention(
5858
# Copy the metadata of the replaced attention node to the new node
5959
# TODO: Investigate why there are multiple FakeTensors in the metadata.
6060
# We only use the first one as it contains the output shape information for this node.
61-
new_attention_node.meta["val"] = copy.copy(
62-
attention_node_replaced.meta["val"][0]
63-
)
61+
if "val" in attention_node_replaced.meta:
62+
new_attention_node.meta["val"] = copy.copy(
63+
attention_node_replaced.meta["val"][0]
64+
)
6465

6566
# If the attention operator had keyword-args, copy them to the new node
6667
if attention_node_replaced.kwargs:

0 commit comments

Comments
 (0)