mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
Fix missing unsqueeze
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user