revert this
This commit is contained in:
@@ -184,9 +184,9 @@ def attention(
|
||||
# )
|
||||
attn_mask = None
|
||||
|
||||
q = q.transpose(1, 2)#.to(dtype)
|
||||
k = k.transpose(1, 2)#.to(dtype)
|
||||
v = v.transpose(1, 2)#.to(dtype)
|
||||
q = q.transpose(1, 2).to(dtype)
|
||||
k = k.transpose(1, 2).to(dtype)
|
||||
v = v.transpose(1, 2).to(dtype)
|
||||
|
||||
out = torch.nn.functional.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
|
||||
@@ -196,9 +196,9 @@ def attention(
|
||||
elif attention_mode == 'sageattn':
|
||||
attn_mask = None
|
||||
|
||||
q = q.transpose(1, 2)#.to(dtype)
|
||||
k = k.transpose(1, 2)#.to(dtype)
|
||||
v = v.transpose(1, 2)#.to(dtype)
|
||||
q = q.transpose(1, 2).to(dtype)
|
||||
k = k.transpose(1, 2).to(dtype)
|
||||
v = v.transpose(1, 2).to(dtype)
|
||||
|
||||
out = sageattn_func(
|
||||
q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
|
||||
|
||||
Reference in New Issue
Block a user