yapf

nod-ai · Apr 24, 2024 · 7f99dfa · 7f99dfa
1 parent 63b7ef8
commit 7f99dfa
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/sharktank/sharktank/models/llama/llama.py b/sharktank/sharktank/models/llama/llama.py
@@ -340,7 +340,9 @@ def forward(
         values = xv.transpose(1, 2)
 
         # Flash attention.
-        attn_output = torch.nn.functional.scaled_dot_product_attention(xq, keys, values, attention_mask)
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            xq, keys, values, attention_mask
+        )
         attn_output = attn_output.transpose(1, 2).reshape(bs, batch_seq_len, -1)
 
         # Project.

diff --git a/sharktank/sharktank/models/llama/llama_ref.py b/sharktank/sharktank/models/llama/llama_ref.py
@@ -209,7 +209,9 @@ def forward(
         values = values.transpose(1, 2)
 
         # Flash attention.
-        attn_output = torch.nn.functional.scaled_dot_product_attention(xq, keys, values, attention_mask)
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            xq, keys, values, attention_mask
+        )
         attn_output = attn_output.transpose(1, 2).reshape(bs, q_len, -1)
 
         # Project.