Fix bug for GaudiMixtralAttentionLongSequence forward (#1650)

Signed-off-by: kaixuanliu <[email protected]>
huggingface · Dec 25, 2024 · 780c9d8 · 780c9d8
1 parent a6ee7c2
commit 780c9d8
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/optimum/habana/transformers/models/mixtral/modeling_mixtral.py b/optimum/habana/transformers/models/mixtral/modeling_mixtral.py
@@ -166,8 +166,7 @@ def forward(q, k, v, mask, causal, q_block_size):
             s, e = i * q_block_size, (i + 1) * q_block_size
             row_q = q[:, :, s:e, :]
             row_mask = mask[:, :, s:e, :]
-            row_o = attn_output[:, :, s:e, :]
-            row_o.fill_(FusedSDPA.apply(row_q, k, v, row_mask, 0.0, causal, None))
+            attn_output[:, :, s:e, :] = FusedSDPA.apply(row_q, k, v, row_mask, 0.0, causal, None)
 
         if q_padding != 0:
             attn_output = attn_output[:, :, :-q_padding, :]