This commit is contained in:
Disty0
2024-01-09 16:07:43 +03:00
parent b206612069
commit 42845b344c
3 changed files with 5 additions and 5 deletions
+2 -2
View File
@@ -149,7 +149,7 @@ class SlicedAttnProcessor: # pylint: disable=too-few-public-methods
hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice
del attn_slice
torch.xpu.synchronize()
torch.xpu.synchronize(query.device)
else:
query_slice = query[start_idx:end_idx]
key_slice = key[start_idx:end_idx]
@@ -284,7 +284,7 @@ class AttnProcessor:
hidden_states[start_idx:end_idx] = attn_slice
del attn_slice
torch.xpu.synchronize()
torch.xpu.synchronize(query.device)
else:
attention_probs = attn.get_attention_scores(query, key, attention_mask)
hidden_states = torch.bmm(attention_probs, value)