chore: fix attn test (#3055)

pytorch · Aug 3, 2024 · 9beac22 · 9beac22
1 parent 411ad4b
commit 9beac22
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/py/torch_tensorrt/dynamo/lowering/passes/lower_scaled_dot_product_attention.py b/py/torch_tensorrt/dynamo/lowering/passes/lower_scaled_dot_product_attention.py
@@ -58,9 +58,10 @@ def lower_scaled_dot_product_attention(
             # Copy the metadata of the replaced attention node to the new node
             # TODO: Investigate why there are multiple FakeTensors in the metadata.
             # We only use the first one as it contains the output shape information for this node.
-            new_attention_node.meta["val"] = copy.copy(
-                attention_node_replaced.meta["val"][0]
-            )
+            if "val" in attention_node_replaced.meta:
+                new_attention_node.meta["val"] = copy.copy(
+                    attention_node_replaced.meta["val"][0]
+                )
 
             # If the attention operator had keyword-args, copy them to the new node
             if attention_node_replaced.kwargs: