Fix missing unsqueeze

This commit is contained in:
Disty0
2025-01-25 21:43:38 +03:00
parent 58a3f7c393
commit 02ecc6710d
2 changed files with 6 additions and 6 deletions
+4 -4
View File
@@ -120,10 +120,10 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop
value[start_idx:end_idx, :, :, :],
attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask,
dropout_p=dropout_p, is_causal=is_causal, **kwargs
)
if is_unsqueezed:
hidden_states.squeeze(0)
)
torch.xpu.synchronize(query.device)
else:
return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
hidden_states = original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
if is_unsqueezed:
hidden_states.squeeze(0)
return hidden_states
+2 -2
View File
@@ -115,12 +115,12 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop
attn_mask=attn_mask[start_idx:end_idx, :, :, :] if attn_mask is not None else attn_mask,
dropout_p=dropout_p, is_causal=is_causal, **kwargs
)
if is_unsqueezed:
hidden_states.squeeze(0)
if devices.backend != "directml":
getattr(torch, query.device.type).synchronize()
else:
return devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
if is_unsqueezed:
hidden_states.squeeze(0)
return hidden_states