Skip to content

GitLab

Commit b373034e authored by Mitchell Wortsman's avatar Mitchell Wortsman
Browse files

test

parent 5f3d9ada
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
tests/triton_tests/attn_decomp.py
View file @ b373034e
...@@ -97,7 +97,7 @@ class Attention(torch.nn.Module): ...@@ -97,7 +97,7 @@ class Attention(torch.nn.Module):
def forward(self, x, attn_mask = None): def forward(self, x, attn_mask = None):
q, k, v = self.in_proj_linear(self.ln(x)).chunk(3, dim=-1) q, k, v = self.in_proj_linear(self.ln(x)).chunk(3, dim=-1)
x = torch.compile(torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask)) x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask)
x = self.out_proj(x) x = self.out_proj(x)
return x return x
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment